From 0cb195ac905f4220e165ead73d4b4f862c77b93d Mon Sep 17 00:00:00 2001 From: Satyam Trivedi Date: Tue, 28 Jul 2026 16:09:10 +0530 Subject: [PATCH 1/5] feat: full AWS Bedrock coverage (#214) Complete the Bedrock SDK-compat surface end-to-end so the real aws-sdk-go-v2 bedrock, bedrockruntime, bedrockagent, and bedrockagentruntime clients all work against the in-memory backend. Runtime (bedrock-runtime): - ConverseStream + InvokeModelWithResponseStream over vnd.amazon.eventstream - CountTokens, ApplyGuardrail - Async invoke: StartAsyncInvoke / GetAsyncInvoke / ListAsyncInvokes Control plane (bedrock): - Guardrail policy configs (topic/content/word/sensitive-info/contextual-grounding) and guardrail versions (CreateGuardrailVersion + version-addressed Get/Delete/List) - Model import jobs, model copy jobs, model evaluation jobs (+ Stop) - Inference profiles, prompt routers - Marketplace model endpoints (incl. Register/Deregister) - Foundation model agreements (Create/Delete/ListOffers/GetAvailability) - Automated reasoning policies (CRUD) - Resource tagging: TagResource / UntagResource / ListTagsForResource (also persists tags previously accepted-but-dropped on guardrail/provisioned create) Agents (new bedrock-agent / bedrock-agent-runtime services + SDK deps): - Agents, knowledge bases, data sources, flows, prompts (CRUD + lifecycle) - Runtime: InvokeAgent (eventstream), Retrieve, RetrieveAndGenerate Implemented across all layers (driver interface, in-memory provider, portable API, restJson1 SDK-compat handler) with unit tests and real-SDK roundtrip tests per subsystem. The two agent services register before the S3 catch-all, and the agent-runtime handler registers before the agent control plane and matches only POST so the shared /agents and /knowledgebases roots never collide. Verified: go build/vet/gofmt clean, golangci-lint clean, go test ./... passing, and a full-server end-to-end exercise with the real SDK clients over a socket. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/sdk-server.md | 12 +- go.mod | 10 +- go.sum | 16 +- providers/aws/aws.go | 104 ++-- providers/aws/bedrock/asyncinvoke.go | 74 +++ .../aws/bedrock/asyncinvoke_jobs_test.go | 178 ++++++ providers/aws/bedrock/bedrock.go | 36 +- .../aws/bedrock/counttokens_applyguardrail.go | 52 ++ .../counttokens_applyguardrail_test.go | 102 +++ providers/aws/bedrock/guardrail_versions.go | 58 ++ providers/aws/bedrock/jobs.go | 262 ++++++++ providers/aws/bedrock/management.go | 122 ++-- providers/aws/bedrock/management_test.go | 107 +++- .../aws/bedrock/marketplace_agreements.go | 215 +++++++ .../bedrock/marketplace_agreements_test.go | 139 +++++ providers/aws/bedrock/registries.go | 287 +++++++++ providers/aws/bedrock/registries_test.go | 181 ++++++ providers/aws/bedrock/tags.go | 110 ++++ providers/aws/bedrock/tags_test.go | 96 +++ providers/aws/bedrockagent/agents.go | 168 +++++ providers/aws/bedrockagent/bedrockagent.go | 53 ++ .../aws/bedrockagent/bedrockagent_test.go | 155 +++++ providers/aws/bedrockagent/datasources.go | 142 +++++ providers/aws/bedrockagent/flows.go | 121 ++++ providers/aws/bedrockagent/knowledgebases.go | 108 ++++ providers/aws/bedrockagent/prompts.go | 99 +++ .../bedrockagentruntime.go | 109 ++++ .../bedrockagentruntime_test.go | 158 +++++ server/aws/aws.go | 106 ++-- server/aws/bedrock/asyncinvoke.go | 210 +++++++ .../aws/bedrock/counttokens_applyguardrail.go | 94 +++ server/aws/bedrock/guardrail_policies.go | 284 +++++++++ server/aws/bedrock/handler.go | 197 +++++- server/aws/bedrock/jobs.go | 517 ++++++++++++++++ server/aws/bedrock/management.go | 117 +++- server/aws/bedrock/marketplace_agreements.go | 400 ++++++++++++ server/aws/bedrock/registries.go | 583 ++++++++++++++++++ .../bedrock/sdk_roundtrip_asyncjobs_test.go | 275 +++++++++ .../sdk_roundtrip_guardrail_policies_test.go | 152 +++++ .../bedrock/sdk_roundtrip_marketplace_test.go | 189 ++++++ .../bedrock/sdk_roundtrip_registries_test.go | 213 +++++++ .../aws/bedrock/sdk_roundtrip_runtime_test.go | 106 ++++ .../bedrock/sdk_roundtrip_streaming_test.go | 118 ++++ server/aws/bedrock/sdk_roundtrip_tags_test.go | 86 +++ server/aws/bedrock/streaming.go | 157 +++++ server/aws/bedrock/tags.go | 77 +++ server/aws/bedrock/types.go | 86 +++ server/aws/bedrockagent/agents.go | 225 +++++++ server/aws/bedrockagent/errors.go | 52 ++ server/aws/bedrockagent/flows.go | 172 ++++++ server/aws/bedrockagent/handler.go | 154 +++++ server/aws/bedrockagent/knowledgebases.go | 346 +++++++++++ server/aws/bedrockagent/prompts.go | 157 +++++ server/aws/bedrockagent/sdk_roundtrip_test.go | 358 +++++++++++ server/aws/bedrockagent/types.go | 289 +++++++++ server/aws/bedrockagentruntime/errors.go | 48 ++ server/aws/bedrockagentruntime/eventstream.go | 55 ++ server/aws/bedrockagentruntime/handler.go | 120 ++++ server/aws/bedrockagentruntime/operations.go | 133 ++++ .../bedrockagentruntime/sdk_roundtrip_test.go | 139 +++++ server/aws/bedrockagentruntime/types.go | 71 +++ services/bedrock/asyncinvoke_jobs.go | 162 +++++ services/bedrock/bedrock.go | 85 ++- services/bedrock/bedrock_test.go | 4 +- services/bedrock/driver/asyncinvoke_jobs.go | 129 ++++ services/bedrock/driver/driver.go | 105 +++- services/bedrock/driver/guardrail_policies.go | 110 ++++ .../bedrock/driver/marketplace_agreements.go | 58 ++ services/bedrock/driver/registries.go | 104 ++++ services/bedrock/marketplace_agreements.go | 146 +++++ services/bedrock/registries.go | 164 +++++ services/bedrockagent/bedrockagent.go | 413 +++++++++++++ services/bedrockagent/driver/driver.go | 238 +++++++ .../bedrockagentruntime.go | 141 +++++ .../bedrockagentruntime_test.go | 66 ++ services/bedrockagentruntime/driver/driver.go | 74 +++ 76 files changed, 11368 insertions(+), 191 deletions(-) create mode 100644 providers/aws/bedrock/asyncinvoke.go create mode 100644 providers/aws/bedrock/asyncinvoke_jobs_test.go create mode 100644 providers/aws/bedrock/counttokens_applyguardrail.go create mode 100644 providers/aws/bedrock/counttokens_applyguardrail_test.go create mode 100644 providers/aws/bedrock/guardrail_versions.go create mode 100644 providers/aws/bedrock/jobs.go create mode 100644 providers/aws/bedrock/marketplace_agreements.go create mode 100644 providers/aws/bedrock/marketplace_agreements_test.go create mode 100644 providers/aws/bedrock/registries.go create mode 100644 providers/aws/bedrock/registries_test.go create mode 100644 providers/aws/bedrock/tags.go create mode 100644 providers/aws/bedrock/tags_test.go create mode 100644 providers/aws/bedrockagent/agents.go create mode 100644 providers/aws/bedrockagent/bedrockagent.go create mode 100644 providers/aws/bedrockagent/bedrockagent_test.go create mode 100644 providers/aws/bedrockagent/datasources.go create mode 100644 providers/aws/bedrockagent/flows.go create mode 100644 providers/aws/bedrockagent/knowledgebases.go create mode 100644 providers/aws/bedrockagent/prompts.go create mode 100644 providers/aws/bedrockagentruntime/bedrockagentruntime.go create mode 100644 providers/aws/bedrockagentruntime/bedrockagentruntime_test.go create mode 100644 server/aws/bedrock/asyncinvoke.go create mode 100644 server/aws/bedrock/counttokens_applyguardrail.go create mode 100644 server/aws/bedrock/guardrail_policies.go create mode 100644 server/aws/bedrock/jobs.go create mode 100644 server/aws/bedrock/marketplace_agreements.go create mode 100644 server/aws/bedrock/registries.go create mode 100644 server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go create mode 100644 server/aws/bedrock/sdk_roundtrip_guardrail_policies_test.go create mode 100644 server/aws/bedrock/sdk_roundtrip_marketplace_test.go create mode 100644 server/aws/bedrock/sdk_roundtrip_registries_test.go create mode 100644 server/aws/bedrock/sdk_roundtrip_runtime_test.go create mode 100644 server/aws/bedrock/sdk_roundtrip_streaming_test.go create mode 100644 server/aws/bedrock/sdk_roundtrip_tags_test.go create mode 100644 server/aws/bedrock/streaming.go create mode 100644 server/aws/bedrock/tags.go create mode 100644 server/aws/bedrockagent/agents.go create mode 100644 server/aws/bedrockagent/errors.go create mode 100644 server/aws/bedrockagent/flows.go create mode 100644 server/aws/bedrockagent/handler.go create mode 100644 server/aws/bedrockagent/knowledgebases.go create mode 100644 server/aws/bedrockagent/prompts.go create mode 100644 server/aws/bedrockagent/sdk_roundtrip_test.go create mode 100644 server/aws/bedrockagent/types.go create mode 100644 server/aws/bedrockagentruntime/errors.go create mode 100644 server/aws/bedrockagentruntime/eventstream.go create mode 100644 server/aws/bedrockagentruntime/handler.go create mode 100644 server/aws/bedrockagentruntime/operations.go create mode 100644 server/aws/bedrockagentruntime/sdk_roundtrip_test.go create mode 100644 server/aws/bedrockagentruntime/types.go create mode 100644 services/bedrock/asyncinvoke_jobs.go create mode 100644 services/bedrock/driver/asyncinvoke_jobs.go create mode 100644 services/bedrock/driver/guardrail_policies.go create mode 100644 services/bedrock/driver/marketplace_agreements.go create mode 100644 services/bedrock/driver/registries.go create mode 100644 services/bedrock/marketplace_agreements.go create mode 100644 services/bedrock/registries.go create mode 100644 services/bedrockagent/bedrockagent.go create mode 100644 services/bedrockagent/driver/driver.go create mode 100644 services/bedrockagentruntime/bedrockagentruntime.go create mode 100644 services/bedrockagentruntime/bedrockagentruntime_test.go create mode 100644 services/bedrockagentruntime/driver/driver.go diff --git a/docs/sdk-server.md b/docs/sdk-server.md index a3d83c92..c362899a 100644 --- a/docs/sdk-server.md +++ b/docs/sdk-server.md @@ -175,7 +175,8 @@ Region, credentials, and tokens can be any dummy values — the server doesn't v | **IAM** *(query protocol)* | Users (Create/Get/List/Delete), Roles (Create/Get/List/Delete), Policies (Create/Get/List/Delete), Attach/Detach/ListAttached for both Users and Roles, Groups (Create/Get/List/Delete + AddUserToGroup/RemoveUserFromGroup/ListGroupsForUser), AccessKeys (Create/List/Delete), InstanceProfiles (Create/Get/List/Delete + AddRoleToInstanceProfile/RemoveRoleFromInstanceProfile). Errors surface as typed `*types.NoSuchEntityException` / `*types.EntityAlreadyExistsException`. | | **Resource Explorer 2** *(JSON)* | Search — free-text plus filter expression over the cross-service inventory; results include ARN, resource type, region, owning account, and tags | | **Resource Groups Tagging API** *(JSON-RPC)* | GetResources (filter by `ResourceTypeFilters` + `TagFilters`, paginated), TagResources, UntagResources, GetTagKeys, GetTagValues | -| **Bedrock** *(REST + JSON, `bedrock` + `bedrock-runtime`)* | Control plane: ListFoundationModels, GetFoundationModel, model-customization jobs (Create/Get/List), custom models (List/Get/Delete), Guardrails (Create/Get/List/Update/Delete), Provisioned Throughput (Create/Get/List/Delete), invocation-logging config (Put/Get/Delete). Runtime: InvokeModel (family-aware response envelopes) and Converse. | +| **Bedrock** *(REST + JSON, `bedrock` + `bedrock-runtime`)* | Control plane: foundation models (List/Get), model-customization jobs (Create/Get/List), custom models (List/Get/Delete), Guardrails (Create/Get/List/Update/Delete + CreateGuardrailVersion, with topic/content/word/sensitive-info/contextual-grounding policy configs and version snapshots), Provisioned Throughput (Create/Get/List/Delete), invocation-logging config (Put/Get/Delete), resource tagging (Tag/Untag/ListTagsForResource), model import jobs, model copy jobs, evaluation jobs (Create/Get/List/Stop), inference profiles (Create/Get/List/Delete), prompt routers (Create/Get/List/Delete), marketplace model endpoints (Create/Get/List/Update/Delete/Register/Deregister), foundation-model agreements (Create/Delete/ListOffers/GetAvailability), automated-reasoning policies (Create/Get/List/Update/Delete). Runtime: InvokeModel (family-aware response envelopes), Converse, ConverseStream + InvokeModelWithResponseStream (eventstream), CountTokens, ApplyGuardrail, and async invoke (Start/Get/List). | +| **Bedrock Agent** *(REST + JSON, `bedrock-agent` + `bedrock-agent-runtime`)* | Control plane: agents (Create/Get/List/Update/Delete/Prepare + alias), knowledge bases (CRUD), data sources (CRUD + StartIngestionJob), flows (CRUD + Prepare), prompts (CRUD). Runtime: InvokeAgent (eventstream), Retrieve, RetrieveAndGenerate. | ### Azure (`server/azure/`) @@ -237,7 +238,9 @@ server/ │ ├── s3/ ec2/ dynamodb/ lambda/ sqs/ cloudwatch/ │ ├── rds/ redshift/ # query-protocol relational DB handlers │ ├── eks/ # REST EKS control-plane handler -│ └── bedrock/ # REST Bedrock control plane + bedrock-runtime +│ ├── bedrock/ # REST Bedrock control plane + bedrock-runtime +│ ├── bedrockagent/ # REST bedrock-agent control plane +│ └── bedrockagentruntime/ # REST bedrock-agent-runtime data plane ├── azure/ │ ├── azure.go # azureserver.New(Drivers{...}) │ ├── virtualmachines/ disks/ snapshots/ images/ sshpublickeys/ @@ -281,7 +284,8 @@ Each handler uses a different signal so dispatch is unambiguous within a provide | AWS Redshift | Form-encoded POST whose `Action=` is a known Redshift operation (registered before EC2) | | AWS EC2 | `Action=…` in URL query or `Content-Type: application/x-www-form-urlencoded` POST | | AWS CloudWatch | `Smithy-Protocol: rpc-v2-cbor` header | -| AWS Bedrock | URL prefix `/foundation-models`, `/model-customization-jobs`, `/custom-models`, `/guardrails`, `/provisioned-model-throughput`, `/logging/modelinvocations`, or bedrock-runtime `/model/{id}/invoke` and `/model/{id}/converse` | +| AWS Bedrock | URL prefix `/foundation-models`, `/model-customization-jobs`, `/custom-models`, `/guardrails`, `/provisioned-model-throughput`, `/logging/modelinvocations`, `/tagResource`, `/untagResource`, `/listTagsForResource`, `/model-import-jobs`, `/model-copy-jobs`, `/evaluation-jobs`, `/evaluation-job/`, `/inference-profiles`, `/prompt-routers`, `/marketplace-model/endpoints`, `/automated-reasoning-policies`, `/create-foundation-model-agreement`, `/delete-foundation-model-agreement`, `/list-foundation-model-agreement-offers/`, `/foundation-model-availability/`, or bedrock-runtime `/model/{id}/{invoke,converse,converse-stream,invoke-with-response-stream,count-tokens}`, `/guardrail/{id}/version/{version}/apply`, and `/async-invoke` | +| AWS Bedrock Agent | Control plane URL prefix `/agents`, `/knowledgebases`, `/flows`, `/prompts`; runtime (registered first, matched only on POST) `/agents/{id}/agentAliases/{a}/sessions/{s}/text` (InvokeAgent), `/knowledgebases/{id}/retrieve` (Retrieve), and `/retrieveAndGenerate` | | AWS S3 | Fallback (everything else REST-shaped) | | Azure (all ARM) | URL begins with `/subscriptions/{sub}` and matches `Microsoft./` | | Azure SQL | ARM provider `Microsoft.Sql` | @@ -319,7 +323,7 @@ Kubernetes ships as **two cooperating handlers**: per-provider control planes (E The data plane intentionally has no controllers — Deployments don't spawn ReplicaSets, Pods stay Pending, Endpoints are empty stubs. RBAC, subresources, PV/PVC, StatefulSet/DaemonSet/Job/CronJob, and Ingress are out of scope. -Two provider-specific services also ship as full SDK-compat handlers. **AWS Bedrock** covers the `bedrock` control plane (foundation models, customization jobs, custom models, guardrails, provisioned throughput, invocation logging) and the `bedrock-runtime` data plane (InvokeModel with family-aware response envelopes, and Converse). **Azure Databricks** covers the `armdatabricks` ARM workspace resource plus the `databricks-sdk-go` workspace data plane — clusters, instance pools, jobs and runs, cluster policies, libraries, permissions, secrets, tokens, git credentials, repos, DBFS, workspace notebooks/directories, SQL warehouses, pipelines, serving endpoints, SCIM identity, and Unity Catalog. +Two provider-specific services also ship as full SDK-compat handlers. **AWS Bedrock** covers the `bedrock` control plane (foundation models, customization jobs, custom models, guardrails with policy configs + versions, provisioned throughput, invocation logging, resource tagging, model import/copy/evaluation jobs, inference profiles, prompt routers, marketplace model endpoints, foundation-model agreements, and automated-reasoning policies) and the `bedrock-runtime` data plane (InvokeModel with family-aware response envelopes, Converse, streaming ConverseStream / InvokeModelWithResponseStream over `vnd.amazon.eventstream`, CountTokens, ApplyGuardrail, and async invoke). A companion **AWS Bedrock Agent** handler covers the `bedrock-agent` control plane (agents, knowledge bases, data sources, flows, prompts) and the `bedrock-agent-runtime` data plane (InvokeAgent streaming, Retrieve, RetrieveAndGenerate); its runtime handler registers before the control plane and matches only POST so the two never collide on the shared `/agents` and `/knowledgebases` roots. **Azure Databricks** covers the `armdatabricks` ARM workspace resource plus the `databricks-sdk-go` workspace data plane — clusters, instance pools, jobs and runs, cluster policies, libraries, permissions, secrets, tokens, git credentials, repos, DBFS, workspace notebooks/directories, SQL warehouses, pipelines, serving endpoints, SCIM identity, and Unity Catalog. The remaining service domains (DNS, Load Balancer, Cache, Secrets, Logging, Notifications, Container Registry, Event Bus) have full driver implementations in `providers/{aws,azure,gcp}/`; SDK-compat handlers are added in lockstep across all 3 providers as each domain ships. diff --git a/go.mod b/go.mod index ef7d57cd..2951ba5b 100644 --- a/go.mod +++ b/go.mod @@ -33,11 +33,14 @@ require ( github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azsecrets v1.5.0 github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.6.4 github.com/Azure/azure-sdk-for-go/sdk/storage/azqueue v1.0.1 - github.com/aws/aws-sdk-go-v2 v1.42.1 + github.com/aws/aws-sdk-go-v2 v1.43.0 + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 github.com/aws/aws-sdk-go-v2/config v1.32.14 github.com/aws/aws-sdk-go-v2/credentials v1.19.14 github.com/aws/aws-sdk-go-v2/service/autoscaling v1.66.1 github.com/aws/aws-sdk-go-v2/service/bedrock v1.64.0 + github.com/aws/aws-sdk-go-v2/service/bedrockagent v1.58.0 + github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime v1.55.0 github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.53.5 github.com/aws/aws-sdk-go-v2/service/cloudwatch v1.56.2 github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs v1.79.0 @@ -92,10 +95,9 @@ require ( github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.32.0 // indirect github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.55.0 // indirect github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 // indirect - github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 // indirect - github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect - github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 // indirect github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 // indirect github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 // indirect github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.9 // indirect diff --git a/go.sum b/go.sum index ece8da11..a3ba1549 100644 --- a/go.sum +++ b/go.sum @@ -108,8 +108,8 @@ github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0 github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.55.0/go.mod h1:vB2GH9GAYYJTO3mEn8oYwzEdhlayZIdQz6zdzgUIRvA= github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 h1:0s6TxfCu2KHkkZPnBfsQ2y5qia0jl3MMrmBhu3nCOYk= github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0/go.mod h1:Mf6O40IAyB9zR/1J8nGDDPirZQQPbYJni8Yisy7NTMc= -github.com/aws/aws-sdk-go-v2 v1.42.1 h1:9eOTgu1z/dVtYpNZ3/8/XbbaX0x/BqE3HUzAzs6K0ek= -github.com/aws/aws-sdk-go-v2 v1.42.1/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= +github.com/aws/aws-sdk-go-v2 v1.43.0 h1:fharf/WhbRAVZ1du0QL7roNFxZ6T/sWr+4Ni617bwSI= +github.com/aws/aws-sdk-go-v2 v1.43.0/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 h1:3IZY0XAJquT3aHzbkHfPzy4ACPcEjVG0x87KOwtpqGY= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14/go.mod h1:zwM6veDkhGgQFqkBy+uT28AAYpLu+uFMlPl+rCg/73E= github.com/aws/aws-sdk-go-v2/config v1.32.14 h1:opVIRo/ZbbI8OIqSOKmpFaY7IwfFUOCCXBsUpJOwDdI= @@ -118,10 +118,10 @@ github.com/aws/aws-sdk-go-v2/credentials v1.19.14 h1:n+UcGWAIZHkXzYt87uMFBv/l8TH github.com/aws/aws-sdk-go-v2/credentials v1.19.14/go.mod h1:cJKuyWB59Mqi0jM3nFYQRmnHVQIcgoxjEMAbLkpr62w= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21 h1:NUS3K4BTDArQqNu2ih7yeDLaS3bmHD0YndtA6UP884g= github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.21/go.mod h1:YWNWJQNjKigKY1RHVJCuupeWDrrHjRqHm0N9rdrWzYI= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 h1:xM/Is9cKMHa8Jj8zkvWhvrFkZsXJV9E+BB4g0HW0duQ= -github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30/go.mod h1:WueJeNDZvK1fMYEWJIkcivBfEzUkTpBhzlrUKKY8EuA= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 h1:jn46zC9LdsVR/ZpMIJqMqb8hHv31BlLx3ulVqNspUOk= -github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30/go.mod h1:1hTMsAgbdS/AtUi4bw8+gUuh1pceo+eXRLfpSuSQj3M= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 h1:Z8F3hfCY33IGpJjFAnv0wvtv1FIKj1GHmRDEYqy64tw= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31/go.mod h1:aVyUoytEyOViR6jhq6jula0xkc5NfBE2hgeF6BvOrao= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 h1:hyOxUyXdh3AyjE93gBgsfziJag9ACwcs+ZpDBLzi8mw= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31/go.mod h1:OERqI9k0draSLB8O8woxY3q25ZWTELRK4RRoLMuMZFo= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6 h1:qYQ4pzQ2Oz6WpQ8T3HvGHnZydA72MnLuFK9tJwmrbHw= github.com/aws/aws-sdk-go-v2/internal/ini v1.8.6/go.mod h1:O3h0IK87yXci+kg6flUKzJnWeziQUKciKrLjcatSNcY= github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.31 h1:3GUprIsfmGcC5SACIyB0e7E0BM1O1b3Erl5CePYIAeQ= @@ -130,6 +130,10 @@ github.com/aws/aws-sdk-go-v2/service/autoscaling v1.66.1 h1:kGlbhb5GMfkP/bcqcbt3 github.com/aws/aws-sdk-go-v2/service/autoscaling v1.66.1/go.mod h1:z45kurrOonQepd3SN5LIgropAn1NGHwBn1yOMF+QVFU= github.com/aws/aws-sdk-go-v2/service/bedrock v1.64.0 h1:jZOg03lM41zl89h17avGr6AFqAb9g3s/rZytTQslz14= github.com/aws/aws-sdk-go-v2/service/bedrock v1.64.0/go.mod h1:f0MnAznWRN75Dnezt6MBnHOKHEpAI/ztr4Q6Lo1IGDk= +github.com/aws/aws-sdk-go-v2/service/bedrockagent v1.58.0 h1:1Yajbk+RDNjb7vXGZ99QMHHGgPd1Ac15KcT34geeZD4= +github.com/aws/aws-sdk-go-v2/service/bedrockagent v1.58.0/go.mod h1:3vsSCey3UL5+d860za3ahqwaF2sv3liH3OwnSdzeyMo= +github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime v1.55.0 h1:6jy8Wc5cScXZvit2sJK3DS6HZ26iUr/lHH8tpwAFlA0= +github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime v1.55.0/go.mod h1:6dK0wIcvx7eeyWpJqiXa4uSuiQcAZqPmE3lWorR4dvA= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.53.5 h1:gWz8Ax1W8BXOoWe3obsIF4ueAY39u9A+NzsY4jppoIw= github.com/aws/aws-sdk-go-v2/service/bedrockruntime v1.53.5/go.mod h1:8m0vIhh44Mmgb+x5o2WzTt0T5NKVtTBhO1j+t7AyvJI= github.com/aws/aws-sdk-go-v2/service/cloudwatch v1.56.2 h1:AEdVlfaKtqjQgnAZ71TAghxd2We92jSez2VAnjOx1vg= diff --git a/providers/aws/aws.go b/providers/aws/aws.go index a9c878b1..3bb41630 100644 --- a/providers/aws/aws.go +++ b/providers/aws/aws.go @@ -5,6 +5,8 @@ import ( "github.com/stackshy/cloudemu/v2/config" "github.com/stackshy/cloudemu/v2/providers/aws/awsiam" "github.com/stackshy/cloudemu/v2/providers/aws/bedrock" + "github.com/stackshy/cloudemu/v2/providers/aws/bedrockagent" + "github.com/stackshy/cloudemu/v2/providers/aws/bedrockagentruntime" "github.com/stackshy/cloudemu/v2/providers/aws/cloudwatch" "github.com/stackshy/cloudemu/v2/providers/aws/cloudwatchlogs" "github.com/stackshy/cloudemu/v2/providers/aws/dynamodb" @@ -30,61 +32,65 @@ import ( // Provider holds all AWS mock services. type Provider struct { - S3 *s3.Mock - EC2 *ec2.Mock - DynamoDB *dynamodb.Mock - Lambda *lambda.Mock - VPC *vpc.Mock - CloudWatch *cloudwatch.Mock - IAM *awsiam.Mock - Route53 *route53.Mock - ELB *elb.Mock - SQS *sqs.Mock - ElastiCache *elasticache.Mock - SecretsManager *secretsmanager.Mock - CloudWatchLogs *cloudwatchlogs.Mock - SNS *sns.Mock - ECR *ecr.Mock - EventBridge *eventbridge.Mock - RDS *rds.Mock - Redshift *redshift.Mock - EKS *eks.Mock - Bedrock *bedrock.Mock - SageMaker *sagemaker.Mock - SSM *ssm.Mock - ResourceDiscovery *resourcediscovery.Engine - AccountID string - Region string + S3 *s3.Mock + EC2 *ec2.Mock + DynamoDB *dynamodb.Mock + Lambda *lambda.Mock + VPC *vpc.Mock + CloudWatch *cloudwatch.Mock + IAM *awsiam.Mock + Route53 *route53.Mock + ELB *elb.Mock + SQS *sqs.Mock + ElastiCache *elasticache.Mock + SecretsManager *secretsmanager.Mock + CloudWatchLogs *cloudwatchlogs.Mock + SNS *sns.Mock + ECR *ecr.Mock + EventBridge *eventbridge.Mock + RDS *rds.Mock + Redshift *redshift.Mock + EKS *eks.Mock + Bedrock *bedrock.Mock + BedrockAgent *bedrockagent.Mock + BedrockAgentRuntime *bedrockagentruntime.Mock + SageMaker *sagemaker.Mock + SSM *ssm.Mock + ResourceDiscovery *resourcediscovery.Engine + AccountID string + Region string } // New creates a new AWS provider with all mock services. func New(opts ...config.Option) *Provider { o := config.NewOptions(opts...) p := &Provider{ - S3: s3.New(o), - EC2: ec2.New(o), - DynamoDB: dynamodb.New(o), - Lambda: lambda.New(o), - VPC: vpc.New(o), - CloudWatch: cloudwatch.New(o), - IAM: awsiam.New(o), - Route53: route53.New(o), - ELB: elb.New(o), - SQS: sqs.New(o), - ElastiCache: elasticache.New(o), - SecretsManager: secretsmanager.New(o), - CloudWatchLogs: cloudwatchlogs.New(o), - SNS: sns.New(o), - ECR: ecr.New(o), - EventBridge: eventbridge.New(o), - RDS: rds.New(o), - Redshift: redshift.New(o), - EKS: eks.New(o), - Bedrock: bedrock.New(o), - SageMaker: sagemaker.New(o), - SSM: ssm.New(o), - AccountID: o.AccountID, - Region: o.Region, + S3: s3.New(o), + EC2: ec2.New(o), + DynamoDB: dynamodb.New(o), + Lambda: lambda.New(o), + VPC: vpc.New(o), + CloudWatch: cloudwatch.New(o), + IAM: awsiam.New(o), + Route53: route53.New(o), + ELB: elb.New(o), + SQS: sqs.New(o), + ElastiCache: elasticache.New(o), + SecretsManager: secretsmanager.New(o), + CloudWatchLogs: cloudwatchlogs.New(o), + SNS: sns.New(o), + ECR: ecr.New(o), + EventBridge: eventbridge.New(o), + RDS: rds.New(o), + Redshift: redshift.New(o), + EKS: eks.New(o), + Bedrock: bedrock.New(o), + BedrockAgent: bedrockagent.New(o), + BedrockAgentRuntime: bedrockagentruntime.New(o), + SageMaker: sagemaker.New(o), + SSM: ssm.New(o), + AccountID: o.AccountID, + Region: o.Region, } p.EC2.SetMonitoring(p.CloudWatch) p.S3.SetMonitoring(p.CloudWatch) diff --git a/providers/aws/bedrock/asyncinvoke.go b/providers/aws/bedrock/asyncinvoke.go new file mode 100644 index 00000000..6193316c --- /dev/null +++ b/providers/aws/bedrock/asyncinvoke.go @@ -0,0 +1,74 @@ +package bedrock + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// StartAsyncInvoke starts an asynchronous model invocation. It completes +// synchronously: the invocation is recorded already in the Completed state so +// Get/List calls are deterministic. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) StartAsyncInvoke(_ context.Context, cfg driver.StartAsyncInvokeConfig) (*driver.AsyncInvoke, error) { + switch { + case cfg.ModelID == "": + return nil, errors.New(errors.InvalidArgument, "modelId is required") + case len(cfg.ModelInput) == 0: + return nil, errors.New(errors.InvalidArgument, "modelInput is required") + case cfg.Output.S3URI == "": + return nil, errors.New(errors.InvalidArgument, "outputDataConfig.s3OutputDataConfig.s3Uri is required") + } + + modelARN := m.resolveModelARN(cfg.ModelID) + if modelARN == "" { + return nil, errors.Newf(errors.InvalidArgument, "model %q not found", cfg.ModelID) + } + + now := m.now() + arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "async-invoke/"+idgen.GenerateID("")) + + inv := &driver.AsyncInvoke{ + InvocationARN: arn, + ModelARN: modelARN, + ClientRequestToken: cfg.ClientRequestToken, + Status: driver.AsyncCompleted, + Output: cfg.Output, + SubmitTime: now, + LastModifiedTime: now, + EndTime: now, + } + m.asyncInvokes.Set(arn, inv) + m.setTags(arn, m.tagsFromMap(cfg.Tags)) + + result := *inv + + return &result, nil +} + +// GetAsyncInvoke returns an async invocation by its invocation ARN. +func (m *Mock) GetAsyncInvoke(_ context.Context, invocationARN string) (*driver.AsyncInvoke, error) { + inv, ok := m.asyncInvokes.Get(invocationARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "async invocation %q not found", invocationARN) + } + + result := *inv + + return &result, nil +} + +// ListAsyncInvokes lists all async invocations. +func (m *Mock) ListAsyncInvokes(_ context.Context) ([]driver.AsyncInvoke, error) { + all := m.asyncInvokes.All() + out := make([]driver.AsyncInvoke, 0, len(all)) + + for _, inv := range all { + out = append(out, *inv) + } + + return out, nil +} diff --git a/providers/aws/bedrock/asyncinvoke_jobs_test.go b/providers/aws/bedrock/asyncinvoke_jobs_test.go new file mode 100644 index 00000000..80fd1bc3 --- /dev/null +++ b/providers/aws/bedrock/asyncinvoke_jobs_test.go @@ -0,0 +1,178 @@ +package bedrock + +import ( + "context" + "testing" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +func TestStartAsyncInvokeLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + inv, err := m.StartAsyncInvoke(ctx, bedrockdriver.StartAsyncInvokeConfig{ + ModelID: titanModel, + ModelInput: []byte(`{"inputText":"hello"}`), + Output: bedrockdriver.AsyncInvokeOutputConfig{S3URI: "s3://bucket/out/"}, + Tags: map[string]string{"team": "ml"}, + }) + requireNoError(t, err) + assertNotEmpty(t, inv.InvocationARN) + assertNotEmpty(t, inv.ModelARN) + assertEqual(t, bedrockdriver.AsyncCompleted, inv.Status) + assertEqual(t, "s3://bucket/out/", inv.Output.S3URI) + + got, err := m.GetAsyncInvoke(ctx, inv.InvocationARN) + requireNoError(t, err) + assertEqual(t, inv.InvocationARN, got.InvocationARN) + assertEqual(t, bedrockdriver.AsyncCompleted, got.Status) + + list, err := m.ListAsyncInvokes(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + tags, err := m.ListTagsForResource(ctx, inv.InvocationARN) + requireNoError(t, err) + assertEqual(t, 1, len(tags)) +} + +func TestStartAsyncInvokeValidation(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.StartAsyncInvoke(ctx, bedrockdriver.StartAsyncInvokeConfig{ + ModelInput: []byte(`{}`), + Output: bedrockdriver.AsyncInvokeOutputConfig{S3URI: "s3://b/o"}, + }) + assertError(t, err, true) + + _, err = m.StartAsyncInvoke(ctx, bedrockdriver.StartAsyncInvokeConfig{ + ModelID: "nope.unknown-v1", + ModelInput: []byte(`{}`), + Output: bedrockdriver.AsyncInvokeOutputConfig{S3URI: "s3://b/o"}, + }) + assertError(t, err, true) +} + +func TestGetAsyncInvokeNotFound(t *testing.T) { + m := newTestMock() + + _, err := m.GetAsyncInvoke(context.Background(), "arn:aws:bedrock:us-east-1:123456789012:async-invoke/missing") + assertError(t, err, true) +} + +func TestModelImportJobLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + job, err := m.CreateModelImportJob(ctx, bedrockdriver.ModelImportJobConfig{ + JobName: "import-1", + ImportedModelName: "my-imported", + RoleARN: "arn:aws:iam::123456789012:role/bedrock", + ModelDataSourceS3URI: "s3://bucket/model/", + }) + requireNoError(t, err) + assertNotEmpty(t, job.JobARN) + assertEqual(t, bedrockdriver.JobCompleted, job.Status) + assertNotEmpty(t, job.ImportedModelARN) + + byName, err := m.GetModelImportJob(ctx, "import-1") + requireNoError(t, err) + assertEqual(t, job.JobARN, byName.JobARN) + + byARN, err := m.GetModelImportJob(ctx, job.JobARN) + requireNoError(t, err) + assertEqual(t, "import-1", byARN.JobName) + + list, err := m.ListModelImportJobs(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + _, err = m.GetModelImportJob(ctx, "missing") + assertError(t, err, true) + + _, err = m.CreateModelImportJob(ctx, bedrockdriver.ModelImportJobConfig{ImportedModelName: "x"}) + assertError(t, err, true) +} + +func TestModelCopyJobLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + src := "arn:aws:bedrock:us-east-1:123456789012:custom-model/src" + job, err := m.CreateModelCopyJob(ctx, bedrockdriver.ModelCopyJobConfig{ + SourceModelARN: src, + TargetModelName: "copy-target", + }) + requireNoError(t, err) + assertNotEmpty(t, job.JobARN) + assertEqual(t, bedrockdriver.JobCompleted, job.Status) + assertEqual(t, "123456789012", job.SourceAccountID) + assertNotEmpty(t, job.TargetModelARN) + + got, err := m.GetModelCopyJob(ctx, job.JobARN) + requireNoError(t, err) + assertEqual(t, src, got.SourceModelARN) + + list, err := m.ListModelCopyJobs(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + _, err = m.GetModelCopyJob(ctx, "arn:aws:bedrock:us-east-1:123456789012:model-copy-job/missing") + assertError(t, err, true) + + _, err = m.CreateModelCopyJob(ctx, bedrockdriver.ModelCopyJobConfig{TargetModelName: "x"}) + assertError(t, err, true) +} + +func TestEvaluationJobLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + job, err := m.CreateEvaluationJob(ctx, bedrockdriver.EvaluationJobConfig{ + JobName: "eval-1", + RoleARN: "arn:aws:iam::123456789012:role/bedrock", + EvaluationConfig: []byte(`{"automated":{}}`), + InferenceConfig: []byte(`{"models":[]}`), + OutputDataS3URI: "s3://bucket/eval/", + }) + requireNoError(t, err) + assertNotEmpty(t, job.JobARN) + assertEqual(t, bedrockdriver.JobCompleted, job.Status) + assertEqual(t, bedrockdriver.EvaluationTypeAutomated, job.JobType) + + got, err := m.GetEvaluationJob(ctx, "eval-1") + requireNoError(t, err) + assertEqual(t, job.JobARN, got.JobARN) + + byARN, err := m.GetEvaluationJob(ctx, job.JobARN) + requireNoError(t, err) + assertEqual(t, "eval-1", byARN.JobName) + + list, err := m.ListEvaluationJobs(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + requireNoError(t, m.StopEvaluationJob(ctx, "eval-1")) + + stopped, err := m.GetEvaluationJob(ctx, "eval-1") + requireNoError(t, err) + assertEqual(t, bedrockdriver.JobStopped, stopped.Status) + + assertError(t, m.StopEvaluationJob(ctx, "missing"), true) +} + +func TestEvaluationJobTypeHuman(t *testing.T) { + m := newTestMock() + + job, err := m.CreateEvaluationJob(context.Background(), bedrockdriver.EvaluationJobConfig{ + JobName: "eval-human", + RoleARN: "arn:aws:iam::123456789012:role/bedrock", + EvaluationConfig: []byte(`{"human":{"humanWorkflowConfig":{}}}`), + InferenceConfig: []byte(`{"models":[]}`), + OutputDataS3URI: "s3://bucket/eval/", + }) + requireNoError(t, err) + assertEqual(t, bedrockdriver.EvaluationTypeHuman, job.JobType) +} diff --git a/providers/aws/bedrock/bedrock.go b/providers/aws/bedrock/bedrock.go index df8c2984..7de3de33 100644 --- a/providers/aws/bedrock/bedrock.go +++ b/providers/aws/bedrock/bedrock.go @@ -24,9 +24,23 @@ type Mock struct { foundation []driver.FoundationModel jobs *memstore.Store[*driver.CustomizationJob] models *memstore.Store[*driver.CustomModel] - guardrails *memstore.Store[*driver.Guardrail] + guardrails *memstore.Store[*guardrailRecord] provisioned *memstore.Store[*driver.ProvisionedThroughput] - opts *config.Options + tags *memstore.Store[[]driver.Tag] // keyed by resource ARN + + asyncInvokes *memstore.Store[*driver.AsyncInvoke] // keyed by invocation ARN + importJobs *memstore.Store[*driver.ModelImportJob] // keyed by job name + copyJobs *memstore.Store[*driver.ModelCopyJob] // keyed by job ARN + evalJobs *memstore.Store[*driver.EvaluationJob] // keyed by job name + + inferenceProfiles *memstore.Store[*driver.InferenceProfile] // keyed by profile ID + promptRouters *memstore.Store[*driver.PromptRouter] // keyed by router ARN + arPolicies *memstore.Store[*driver.AutomatedReasoningPolicy] // keyed by policy ARN + + marketplaceEndpoints *memstore.Store[*driver.MarketplaceEndpoint] // keyed by endpoint ARN + fmAgreements *memstore.Store[bool] // set of accepted agreements keyed by modelId + + opts *config.Options logMu sync.RWMutex logging *driver.LoggingConfig @@ -39,9 +53,23 @@ func New(opts *config.Options) *Mock { foundation: seedFoundationModels(opts.Region), jobs: memstore.New[*driver.CustomizationJob](), models: memstore.New[*driver.CustomModel](), - guardrails: memstore.New[*driver.Guardrail](), + guardrails: memstore.New[*guardrailRecord](), provisioned: memstore.New[*driver.ProvisionedThroughput](), - opts: opts, + tags: memstore.New[[]driver.Tag](), + + asyncInvokes: memstore.New[*driver.AsyncInvoke](), + importJobs: memstore.New[*driver.ModelImportJob](), + copyJobs: memstore.New[*driver.ModelCopyJob](), + evalJobs: memstore.New[*driver.EvaluationJob](), + + inferenceProfiles: memstore.New[*driver.InferenceProfile](), + promptRouters: memstore.New[*driver.PromptRouter](), + arPolicies: memstore.New[*driver.AutomatedReasoningPolicy](), + + marketplaceEndpoints: memstore.New[*driver.MarketplaceEndpoint](), + fmAgreements: memstore.New[bool](), + + opts: opts, } } diff --git a/providers/aws/bedrock/counttokens_applyguardrail.go b/providers/aws/bedrock/counttokens_applyguardrail.go new file mode 100644 index 00000000..6fd2c3b0 --- /dev/null +++ b/providers/aws/bedrock/counttokens_applyguardrail.go @@ -0,0 +1,52 @@ +package bedrock + +import ( + "context" + "strings" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// CountTokens estimates the input token count for a would-be inference request. +// When a raw InvokeModel body is supplied it counts the extracted prompt; +// otherwise it counts the Converse messages plus any system content. +// +//nolint:gocritic // in matches the driver interface signature; read without mutation. +func (m *Mock) CountTokens(_ context.Context, in driver.CountTokensInput) (int, error) { + if in.ModelID == "" { + return 0, errors.New(errors.InvalidArgument, "modelId is required") + } + + if !m.modelExists(in.ModelID) { + return 0, errors.Newf(errors.InvalidArgument, "model %q not found", in.ModelID) + } + + if len(in.InvokeBody) > 0 { + return wordCount(extractPrompt(in.InvokeBody)), nil + } + + return conversationTokens(in.Messages) + wordCount(strings.Join(in.System, " ")), nil +} + +// ApplyGuardrail evaluates content against a guardrail. The emulator never +// intervenes: it validates the request and echoes the input content back +// through as outputs with a NONE action. +func (m *Mock) ApplyGuardrail(_ context.Context, in driver.ApplyGuardrailInput) (*driver.ApplyGuardrailOutput, error) { + if in.GuardrailIdentifier == "" { + return nil, errors.New(errors.InvalidArgument, "guardrailIdentifier is required") + } + + if m.findGuardrailRecord(in.GuardrailIdentifier) == nil { + return nil, errors.Newf(errors.NotFound, "guardrail %q not found", in.GuardrailIdentifier) + } + + if in.Source != driver.GuardrailSourceInput && in.Source != driver.GuardrailSourceOutput { + return nil, errors.Newf(errors.InvalidArgument, "invalid source %q: want INPUT or OUTPUT", in.Source) + } + + return &driver.ApplyGuardrailOutput{ + Action: driver.GuardrailActionNone, + Outputs: append([]string(nil), in.Content...), + }, nil +} diff --git a/providers/aws/bedrock/counttokens_applyguardrail_test.go b/providers/aws/bedrock/counttokens_applyguardrail_test.go new file mode 100644 index 00000000..e1ea4543 --- /dev/null +++ b/providers/aws/bedrock/counttokens_applyguardrail_test.go @@ -0,0 +1,102 @@ +package bedrock + +import ( + "context" + "testing" + + "github.com/stackshy/cloudemu/v2/errors" + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +func TestCountTokensConverse(t *testing.T) { + m := newTestMock() + + n, err := m.CountTokens(context.Background(), bedrockdriver.CountTokensInput{ + ModelID: titanModel, + System: []string{"Be concise."}, + Messages: []bedrockdriver.Message{{Role: "user", Text: []string{"What is Bedrock?"}}}, + }) + requireNoError(t, err) + + if n <= 0 { + t.Fatalf("expected a positive token count, got %d", n) + } +} + +func TestCountTokensInvokeBody(t *testing.T) { + m := newTestMock() + + n, err := m.CountTokens(context.Background(), bedrockdriver.CountTokensInput{ + ModelID: titanModel, + InvokeBody: []byte(`{"inputText":"hello there world"}`), + }) + requireNoError(t, err) + assertEqual(t, 3, n) +} + +func TestCountTokensInvalidModel(t *testing.T) { + m := newTestMock() + + _, err := m.CountTokens(context.Background(), bedrockdriver.CountTokensInput{ + ModelID: "nope.unknown-v1", + Messages: []bedrockdriver.Message{{Role: "user", Text: []string{"hi"}}}, + }) + assertError(t, err, true) + + _, err = m.CountTokens(context.Background(), bedrockdriver.CountTokensInput{}) + assertError(t, err, true) +} + +func TestApplyGuardrailEchoes(t *testing.T) { + m := newTestMock() + g := newGuardrail(t, m, "gr-apply") + + out, err := m.ApplyGuardrail(context.Background(), bedrockdriver.ApplyGuardrailInput{ + GuardrailIdentifier: g.ID, + GuardrailVersion: "DRAFT", + Source: bedrockdriver.GuardrailSourceInput, + Content: []string{"hello", "world"}, + }) + requireNoError(t, err) + assertEqual(t, bedrockdriver.GuardrailActionNone, out.Action) + assertEqual(t, 2, len(out.Outputs)) + assertEqual(t, "hello", out.Outputs[0]) + assertEqual(t, "world", out.Outputs[1]) +} + +func TestApplyGuardrailUnknown(t *testing.T) { + m := newTestMock() + + _, err := m.ApplyGuardrail(context.Background(), bedrockdriver.ApplyGuardrailInput{ + GuardrailIdentifier: "gr-missing", + Source: bedrockdriver.GuardrailSourceInput, + Content: []string{"hi"}, + }) + assertError(t, err, true) + + if !errors.IsNotFound(err) { + t.Fatalf("expected NotFound error, got %v", err) + } +} + +func TestApplyGuardrailBadSource(t *testing.T) { + m := newTestMock() + g := newGuardrail(t, m, "gr-src") + + _, err := m.ApplyGuardrail(context.Background(), bedrockdriver.ApplyGuardrailInput{ + GuardrailIdentifier: g.ID, + Source: "SIDEWAYS", + Content: []string{"hi"}, + }) + assertError(t, err, true) + + if !errors.IsInvalidArgument(err) { + t.Fatalf("expected InvalidArgument error, got %v", err) + } + + _, err = m.ApplyGuardrail(context.Background(), bedrockdriver.ApplyGuardrailInput{ + GuardrailIdentifier: "", + Source: bedrockdriver.GuardrailSourceInput, + }) + assertError(t, err, true) +} diff --git a/providers/aws/bedrock/guardrail_versions.go b/providers/aws/bedrock/guardrail_versions.go new file mode 100644 index 00000000..a0a9ad44 --- /dev/null +++ b/providers/aws/bedrock/guardrail_versions.go @@ -0,0 +1,58 @@ +package bedrock + +import ( + "context" + "strconv" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// guardrailRecord holds a guardrail's mutable working copy (the "DRAFT" +// version) alongside its immutable, numbered version snapshots. Records are +// stored in m.guardrails keyed by the guardrail name. +type guardrailRecord struct { + draft *driver.Guardrail + versions []*driver.Guardrail // numbered snapshots ("1", "2", ...) in creation order + nextVer int // monotonic next version number; never reused after delete +} + +// version returns the snapshot for the given numbered version, or nil. +func (rec *guardrailRecord) version(v string) *driver.Guardrail { + for _, g := range rec.versions { + if g.Version == v { + return g + } + } + + return nil +} + +// CreateGuardrailVersion snapshots the current DRAFT into a new immutable, +// numbered version and returns the guardrail ID and the assigned version. +func (m *Mock) CreateGuardrailVersion( + _ context.Context, identifier, description string, +) (guardrailID, version string, err error) { + rec := m.findGuardrailRecord(identifier) + if rec == nil { + return "", "", errors.Newf(errors.NotFound, "guardrail %q not found", identifier) + } + + ver := strconv.Itoa(rec.nextVer) + rec.nextVer++ + + snapshot := *rec.draft + snapshot.Version = ver + + if description != "" { + snapshot.Description = description + } + + now := m.now() + snapshot.CreatedAt = now + snapshot.UpdatedAt = now + + rec.versions = append(rec.versions, &snapshot) + + return snapshot.ID, ver, nil +} diff --git a/providers/aws/bedrock/jobs.go b/providers/aws/bedrock/jobs.go new file mode 100644 index 00000000..46414deb --- /dev/null +++ b/providers/aws/bedrock/jobs.go @@ -0,0 +1,262 @@ +package bedrock + +import ( + "context" + "encoding/json" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- Model import jobs --- + +// CreateModelImportJob starts a custom-model import job. It completes +// synchronously: the job is recorded already Completed with a materialized +// imported-model ARN. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateModelImportJob(_ context.Context, cfg driver.ModelImportJobConfig) (*driver.ModelImportJob, error) { + switch { + case cfg.JobName == "": + return nil, errors.New(errors.InvalidArgument, "jobName is required") + case cfg.ImportedModelName == "": + return nil, errors.New(errors.InvalidArgument, "importedModelName is required") + case cfg.RoleARN == "": + return nil, errors.New(errors.InvalidArgument, "roleArn is required") + case cfg.ModelDataSourceS3URI == "": + return nil, errors.New(errors.InvalidArgument, "modelDataSource.s3DataSource.s3Uri is required") + } + + if m.importJobs.Has(cfg.JobName) { + return nil, errors.Newf(errors.AlreadyExists, "model import job %q already exists", cfg.JobName) + } + + now := m.now() + jobARN := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "model-import-job/"+idgen.GenerateID("")) + modelARN := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "imported-model/"+cfg.ImportedModelName) + + job := &driver.ModelImportJob{ + JobARN: jobARN, + JobName: cfg.JobName, + ImportedModelName: cfg.ImportedModelName, + ImportedModelARN: modelARN, + RoleARN: cfg.RoleARN, + ModelDataSourceS3URI: cfg.ModelDataSourceS3URI, + Status: driver.JobCompleted, + CreationTime: now, + LastModifiedTime: now, + EndTime: now, + } + m.importJobs.Set(cfg.JobName, job) + m.setTags(jobARN, m.tagsFromMap(cfg.JobTags)) + m.setTags(modelARN, m.tagsFromMap(cfg.ImportedModelTags)) + + result := *job + + return &result, nil +} + +// GetModelImportJob returns an import job by name or ARN. +func (m *Mock) GetModelImportJob(_ context.Context, jobIdentifier string) (*driver.ModelImportJob, error) { + if job, ok := m.importJobs.Get(jobIdentifier); ok { + result := *job + + return &result, nil + } + + for _, job := range m.importJobs.All() { + if job.JobARN == jobIdentifier { + result := *job + + return &result, nil + } + } + + return nil, errors.Newf(errors.NotFound, "model import job %q not found", jobIdentifier) +} + +// ListModelImportJobs lists all import jobs. +func (m *Mock) ListModelImportJobs(_ context.Context) ([]driver.ModelImportJob, error) { + all := m.importJobs.All() + out := make([]driver.ModelImportJob, 0, len(all)) + + for _, job := range all { + out = append(out, *job) + } + + return out, nil +} + +// --- Model copy jobs --- + +// CreateModelCopyJob starts a model-copy job. It completes synchronously: the +// job is recorded already Completed with a materialized target-model ARN. +func (m *Mock) CreateModelCopyJob(_ context.Context, cfg driver.ModelCopyJobConfig) (*driver.ModelCopyJob, error) { + switch { + case cfg.SourceModelARN == "": + return nil, errors.New(errors.InvalidArgument, "sourceModelArn is required") + case cfg.TargetModelName == "": + return nil, errors.New(errors.InvalidArgument, "targetModelName is required") + } + + now := m.now() + jobARN := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "model-copy-job/"+idgen.GenerateID("")) + targetARN := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "model-copy-target/"+cfg.TargetModelName) + + job := &driver.ModelCopyJob{ + JobARN: jobARN, + SourceAccountID: m.opts.AccountID, + SourceModelARN: cfg.SourceModelARN, + TargetModelName: cfg.TargetModelName, + TargetModelARN: targetARN, + TargetModelKMSKeyARN: cfg.ModelKMSKeyID, + Status: driver.JobCompleted, + CreationTime: now, + } + m.copyJobs.Set(jobARN, job) + m.setTags(targetARN, m.tagsFromMap(cfg.TargetModelTags)) + + result := *job + + return &result, nil +} + +// GetModelCopyJob returns a copy job by its job ARN. +func (m *Mock) GetModelCopyJob(_ context.Context, jobARN string) (*driver.ModelCopyJob, error) { + job, ok := m.copyJobs.Get(jobARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "model copy job %q not found", jobARN) + } + + result := *job + + return &result, nil +} + +// ListModelCopyJobs lists all copy jobs. +func (m *Mock) ListModelCopyJobs(_ context.Context) ([]driver.ModelCopyJob, error) { + all := m.copyJobs.All() + out := make([]driver.ModelCopyJob, 0, len(all)) + + for _, job := range all { + out = append(out, *job) + } + + return out, nil +} + +// --- Evaluation jobs --- + +// CreateEvaluationJob starts a model-evaluation job. It completes +// synchronously: the job is recorded already Completed. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateEvaluationJob(_ context.Context, cfg driver.EvaluationJobConfig) (*driver.EvaluationJob, error) { + switch { + case cfg.JobName == "": + return nil, errors.New(errors.InvalidArgument, "jobName is required") + case cfg.RoleARN == "": + return nil, errors.New(errors.InvalidArgument, "roleArn is required") + case len(cfg.EvaluationConfig) == 0: + return nil, errors.New(errors.InvalidArgument, "evaluationConfig is required") + case len(cfg.InferenceConfig) == 0: + return nil, errors.New(errors.InvalidArgument, "inferenceConfig is required") + case cfg.OutputDataS3URI == "": + return nil, errors.New(errors.InvalidArgument, "outputDataConfig.s3Uri is required") + } + + if m.evalJobs.Has(cfg.JobName) { + return nil, errors.Newf(errors.AlreadyExists, "evaluation job %q already exists", cfg.JobName) + } + + now := m.now() + jobARN := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "evaluation-job/"+idgen.GenerateID("")) + + job := &driver.EvaluationJob{ + JobARN: jobARN, + JobName: cfg.JobName, + JobType: evaluationJobType(cfg.EvaluationConfig), + ApplicationType: cfg.ApplicationType, + RoleARN: cfg.RoleARN, + EvaluationConfig: cfg.EvaluationConfig, + InferenceConfig: cfg.InferenceConfig, + OutputDataS3URI: cfg.OutputDataS3URI, + JobDescription: cfg.JobDescription, + CustomerEncryptionKeyID: cfg.CustomerEncryptionKeyID, + Status: driver.JobCompleted, + CreationTime: now, + LastModifiedTime: now, + } + m.evalJobs.Set(cfg.JobName, job) + m.setTags(jobARN, m.tagsFromMap(cfg.JobTags)) + + result := *job + + return &result, nil +} + +// GetEvaluationJob returns an evaluation job by name or ARN. +func (m *Mock) GetEvaluationJob(_ context.Context, jobIdentifier string) (*driver.EvaluationJob, error) { + job := m.findEvalJob(jobIdentifier) + if job == nil { + return nil, errors.Newf(errors.NotFound, "evaluation job %q not found", jobIdentifier) + } + + result := *job + + return &result, nil +} + +// ListEvaluationJobs lists all evaluation jobs. +func (m *Mock) ListEvaluationJobs(_ context.Context) ([]driver.EvaluationJob, error) { + all := m.evalJobs.All() + out := make([]driver.EvaluationJob, 0, len(all)) + + for _, job := range all { + out = append(out, *job) + } + + return out, nil +} + +// StopEvaluationJob transitions an evaluation job to the Stopped state. +func (m *Mock) StopEvaluationJob(_ context.Context, jobIdentifier string) error { + job := m.findEvalJob(jobIdentifier) + if job == nil { + return errors.Newf(errors.NotFound, "evaluation job %q not found", jobIdentifier) + } + + job.Status = driver.JobStopped + job.LastModifiedTime = m.now() + m.evalJobs.Set(job.JobName, job) + + return nil +} + +func (m *Mock) findEvalJob(id string) *driver.EvaluationJob { + if job, ok := m.evalJobs.Get(id); ok { + return job + } + + for _, job := range m.evalJobs.All() { + if job.JobARN == id { + return job + } + } + + return nil +} + +// evaluationJobType derives the job type from the evaluationConfig document: a +// "human" member yields Human, otherwise Automated. +func evaluationJobType(cfg []byte) string { + var probe map[string]json.RawMessage + if json.Unmarshal(cfg, &probe) == nil { + if _, ok := probe["human"]; ok { + return driver.EvaluationTypeHuman + } + } + + return driver.EvaluationTypeAutomated +} diff --git a/providers/aws/bedrock/management.go b/providers/aws/bedrock/management.go index 4650bc54..8309b4f0 100644 --- a/providers/aws/bedrock/management.go +++ b/providers/aws/bedrock/management.go @@ -39,93 +39,136 @@ func (m *Mock) CreateGuardrail(_ context.Context, cfg driver.GuardrailConfig) (* BlockedOutputsMessaging: cfg.BlockedOutputsMessaging, CreatedAt: now, UpdatedAt: now, + GuardrailPolicies: cfg.GuardrailPolicies, } if cfg.KMSKeyID != "" { g.KMSKeyARN = cfg.KMSKeyID } - m.guardrails.Set(cfg.Name, g) + m.guardrails.Set(cfg.Name, &guardrailRecord{draft: g, nextVer: 1}) + m.setTags(g.ARN, m.tagsFromMap(cfg.Tags)) result := *g return &result, nil } -// GetGuardrail returns a guardrail by name or ARN. version is ignored (only -// the working version is modeled). -func (m *Mock) GetGuardrail(_ context.Context, identifier, _ string) (*driver.Guardrail, error) { - g := m.findGuardrail(identifier) - if g == nil { +// GetGuardrail returns a guardrail by name or ARN. An empty or "DRAFT" version +// returns the working copy; a numbered version returns that snapshot. +func (m *Mock) GetGuardrail(_ context.Context, identifier, version string) (*driver.Guardrail, error) { + rec := m.findGuardrailRecord(identifier) + if rec == nil { return nil, errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } + g := rec.draft + if version != "" && version != guardrailDraftVersion { + if g = rec.version(version); g == nil { + return nil, errors.Newf(errors.NotFound, "guardrail %q version %q not found", identifier, version) + } + } + result := *g return &result, nil } -// ListGuardrails lists all guardrails. -func (m *Mock) ListGuardrails(_ context.Context) ([]driver.Guardrail, error) { +// ListGuardrails lists guardrails. When identifier is set, it returns one +// summary per version (DRAFT plus each numbered snapshot) of that guardrail; +// otherwise it returns one summary per guardrail (its DRAFT). +func (m *Mock) ListGuardrails(_ context.Context, identifier string) ([]driver.Guardrail, error) { + if identifier != "" { + rec := m.findGuardrailRecord(identifier) + if rec == nil { + return nil, errors.Newf(errors.NotFound, "guardrail %q not found", identifier) + } + + out := make([]driver.Guardrail, 0, len(rec.versions)+1) + out = append(out, *rec.draft) + + for _, v := range rec.versions { + out = append(out, *v) + } + + return out, nil + } + all := m.guardrails.All() out := make([]driver.Guardrail, 0, len(all)) - for _, g := range all { - out = append(out, *g) + for _, rec := range all { + out = append(out, *rec.draft) } return out, nil } -// UpdateGuardrail updates a guardrail's mutable fields. +// UpdateGuardrail updates a guardrail's mutable DRAFT working copy. Numbered +// version snapshots are immutable and left untouched. // //nolint:gocritic // cfg matches the driver interface signature; copied once on entry. func (m *Mock) UpdateGuardrail(_ context.Context, identifier string, cfg driver.GuardrailConfig) (*driver.Guardrail, error) { - g := m.findGuardrail(identifier) - if g == nil { + rec := m.findGuardrailRecord(identifier) + if rec == nil { return nil, errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } - updated := *g - updated.Name = orDefault(cfg.Name, g.Name) - updated.Description = cfg.Description - updated.BlockedInputMessaging = orDefault(cfg.BlockedInputMessaging, g.BlockedInputMessaging) - updated.BlockedOutputsMessaging = orDefault(cfg.BlockedOutputsMessaging, g.BlockedOutputsMessaging) - updated.UpdatedAt = m.now() - - // Guardrails are keyed by name; re-key when an update renames one so - // lookups by the new name keep working. - if updated.Name != g.Name { - m.guardrails.Delete(g.Name) + g := rec.draft + oldName := g.Name + g.Name = orDefault(cfg.Name, g.Name) + g.Description = cfg.Description + g.BlockedInputMessaging = orDefault(cfg.BlockedInputMessaging, g.BlockedInputMessaging) + g.BlockedOutputsMessaging = orDefault(cfg.BlockedOutputsMessaging, g.BlockedOutputsMessaging) + g.GuardrailPolicies = cfg.GuardrailPolicies + g.UpdatedAt = m.now() + + // Records are keyed by name; re-key when an update renames one so lookups + // by the new name keep working. + if g.Name != oldName { + m.guardrails.Delete(oldName) + m.guardrails.Set(g.Name, rec) } - m.guardrails.Set(updated.Name, &updated) - - result := updated + result := *g return &result, nil } -// DeleteGuardrail deletes a guardrail by name or ARN. -func (m *Mock) DeleteGuardrail(_ context.Context, identifier string) error { - g := m.findGuardrail(identifier) - if g == nil { +// DeleteGuardrail deletes a guardrail by name or ARN. An empty version deletes +// the whole guardrail (DRAFT and all snapshots); a specific version deletes +// just that numbered snapshot. +func (m *Mock) DeleteGuardrail(_ context.Context, identifier, version string) error { + rec := m.findGuardrailRecord(identifier) + if rec == nil { return errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } - m.guardrails.Delete(g.Name) + if version == "" { + m.guardrails.Delete(rec.draft.Name) - return nil + return nil + } + + for i, v := range rec.versions { + if v.Version == version { + rec.versions = append(rec.versions[:i], rec.versions[i+1:]...) + + return nil + } + } + + return errors.Newf(errors.NotFound, "guardrail %q version %q not found", identifier, version) } -func (m *Mock) findGuardrail(id string) *driver.Guardrail { - if g, ok := m.guardrails.Get(id); ok { - return g +func (m *Mock) findGuardrailRecord(id string) *guardrailRecord { + if rec, ok := m.guardrails.Get(id); ok { + return rec } - for _, g := range m.guardrails.All() { - if g.ID == id || g.ARN == id { - return g + for _, rec := range m.guardrails.All() { + if rec.draft.ID == id || rec.draft.ARN == id { + return rec } } @@ -187,6 +230,7 @@ func (m *Mock) CreateProvisionedModelThroughput( LastModifiedTime: now, } m.provisioned.Set(cfg.ProvisionedModelName, pt) + m.setTags(pt.ARN, m.tagsFromMap(cfg.Tags)) result := *pt diff --git a/providers/aws/bedrock/management_test.go b/providers/aws/bedrock/management_test.go index 530ef6c0..f5657d99 100644 --- a/providers/aws/bedrock/management_test.go +++ b/providers/aws/bedrock/management_test.go @@ -37,7 +37,7 @@ func TestGuardrailLifecycle(t *testing.T) { _, err = m.GetGuardrail(ctx, g.ARN, "") requireNoError(t, err) - list, err := m.ListGuardrails(ctx) + list, err := m.ListGuardrails(ctx, "") requireNoError(t, err) assertEqual(t, 1, len(list)) @@ -49,7 +49,7 @@ func TestGuardrailLifecycle(t *testing.T) { requireNoError(t, err) assertEqual(t, "new in", upd.BlockedInputMessaging) - requireNoError(t, m.DeleteGuardrail(ctx, g.ID)) + requireNoError(t, m.DeleteGuardrail(ctx, g.ID, "")) _, err = m.GetGuardrail(ctx, g.ID, "") assertError(t, err, true) @@ -76,7 +76,7 @@ func TestGuardrailRenameRekeys(t *testing.T) { _, err = m.GetGuardrail(ctx, "gr-old", "") assertError(t, err, true) - list, err := m.ListGuardrails(ctx) + list, err := m.ListGuardrails(ctx, "") requireNoError(t, err) assertEqual(t, 1, len(list)) } @@ -98,6 +98,107 @@ func TestGuardrailValidation(t *testing.T) { assertError(t, err, true) } +func TestGuardrailPoliciesRoundTrip(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateGuardrail(ctx, bedrockdriver.GuardrailConfig{ + Name: "gr-pol", + BlockedInputMessaging: "in", + BlockedOutputsMessaging: "out", + GuardrailPolicies: bedrockdriver.GuardrailPolicies{ + TopicPolicy: &bedrockdriver.GuardrailTopicPolicy{Topics: []bedrockdriver.GuardrailTopic{ + {Name: "fiduciary", Definition: "financial advice", Examples: []string{"invest?"}, Type: "DENY"}, + }}, + ContentPolicy: &bedrockdriver.GuardrailContentPolicy{Filters: []bedrockdriver.GuardrailContentFilter{ + {Type: "HATE", InputStrength: "HIGH", OutputStrength: "MEDIUM"}, + }}, + SensitiveInformationPolicy: &bedrockdriver.GuardrailSensitiveInformationPolicy{ + PiiEntities: []bedrockdriver.GuardrailPiiEntity{{Type: "EMAIL", Action: "ANONYMIZE"}}, + }, + ContextualGroundingPolicy: &bedrockdriver.GuardrailContextualGroundingPolicy{ + Filters: []bedrockdriver.GuardrailContextualGroundingFilter{{Type: "GROUNDING", Threshold: 0.7, Action: "BLOCK"}}, + }, + }, + }) + requireNoError(t, err) + + got, err := m.GetGuardrail(ctx, "gr-pol", "") + requireNoError(t, err) + + if got.TopicPolicy == nil || len(got.TopicPolicy.Topics) != 1 || got.TopicPolicy.Topics[0].Name != "fiduciary" { + t.Fatalf("topic policy not round-tripped: %+v", got.TopicPolicy) + } + + if got.ContentPolicy == nil || got.ContentPolicy.Filters[0].InputStrength != "HIGH" { + t.Fatalf("content policy not round-tripped: %+v", got.ContentPolicy) + } + + if got.SensitiveInformationPolicy == nil || got.SensitiveInformationPolicy.PiiEntities[0].Action != "ANONYMIZE" { + t.Fatalf("sensitive-info policy not round-tripped: %+v", got.SensitiveInformationPolicy) + } + + if got.ContextualGroundingPolicy == nil || got.ContextualGroundingPolicy.Filters[0].Threshold != 0.7 { + t.Fatalf("contextual-grounding policy not round-tripped: %+v", got.ContextualGroundingPolicy) + } +} + +func TestGuardrailVersions(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + g := newGuardrail(t, m, "gr-ver") + + gid, v1, err := m.CreateGuardrailVersion(ctx, g.ID, "first snapshot") + requireNoError(t, err) + assertEqual(t, g.ID, gid) + assertEqual(t, "1", v1) + + // A second version increments. + _, v2, err := m.CreateGuardrailVersion(ctx, g.ID, "") + requireNoError(t, err) + assertEqual(t, "2", v2) + + // The numbered snapshot is retrievable and carries its own description. + snap, err := m.GetGuardrail(ctx, g.ID, "1") + requireNoError(t, err) + assertEqual(t, "1", snap.Version) + assertEqual(t, "first snapshot", snap.Description) + + // DRAFT still resolves to the working copy. + draft, err := m.GetGuardrail(ctx, g.ID, "DRAFT") + requireNoError(t, err) + assertEqual(t, "DRAFT", draft.Version) + + // Unknown version → NotFound. + _, err = m.GetGuardrail(ctx, g.ID, "99") + assertError(t, err, true) + + // Scoped list returns DRAFT + both numbered versions. + scoped, err := m.ListGuardrails(ctx, g.ID) + requireNoError(t, err) + assertEqual(t, 3, len(scoped)) + + // Unscoped list returns one entry per guardrail. + all, err := m.ListGuardrails(ctx, "") + requireNoError(t, err) + assertEqual(t, 1, len(all)) + + // Deleting a specific version leaves DRAFT and the other version intact. + requireNoError(t, m.DeleteGuardrail(ctx, g.ID, "1")) + _, err = m.GetGuardrail(ctx, g.ID, "1") + assertError(t, err, true) + + scoped, err = m.ListGuardrails(ctx, g.ID) + requireNoError(t, err) + assertEqual(t, 2, len(scoped)) + + // Deleting with no version removes the whole guardrail. + requireNoError(t, m.DeleteGuardrail(ctx, g.ID, "")) + _, err = m.GetGuardrail(ctx, g.ID, "") + assertError(t, err, true) +} + func TestProvisionedThroughputLifecycle(t *testing.T) { m := newTestMock() ctx := context.Background() diff --git a/providers/aws/bedrock/marketplace_agreements.go b/providers/aws/bedrock/marketplace_agreements.go new file mode 100644 index 00000000..c280da76 --- /dev/null +++ b/providers/aws/bedrock/marketplace_agreements.go @@ -0,0 +1,215 @@ +package bedrock + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- Marketplace model endpoints --- + +// CreateMarketplaceModelEndpoint deploys a marketplace model endpoint. The +// endpoint ARN is a SageMaker endpoint ARN and the endpoint is REGISTERED and +// InService immediately. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateMarketplaceModelEndpoint( + _ context.Context, cfg driver.MarketplaceEndpointConfig, +) (*driver.MarketplaceEndpoint, error) { + switch { + case cfg.EndpointName == "": + return nil, errors.New(errors.InvalidArgument, "endpointName is required") + case cfg.ModelSourceIdentifier == "": + return nil, errors.New(errors.InvalidArgument, "modelSourceIdentifier is required") + case len(cfg.EndpointConfig) == 0: + return nil, errors.New(errors.InvalidArgument, "endpointConfig is required") + } + + now := m.now() + arn := idgen.AWSARN("sagemaker", m.opts.Region, m.opts.AccountID, "endpoint/"+cfg.EndpointName) + + endpoint := &driver.MarketplaceEndpoint{ + EndpointARN: arn, + ModelSourceIdentifier: cfg.ModelSourceIdentifier, + EndpointConfig: copyBytes(cfg.EndpointConfig), + EndpointStatus: driver.MarketplaceEndpointStatusInService, + Status: driver.MarketplaceEndpointStatusRegistered, + CreatedAt: now, + UpdatedAt: now, + } + m.marketplaceEndpoints.Set(arn, endpoint) + m.setTags(arn, m.tagsFromMap(cfg.Tags)) + + result := *endpoint + + return &result, nil +} + +// GetMarketplaceModelEndpoint returns a marketplace model endpoint by ARN. +func (m *Mock) GetMarketplaceModelEndpoint(_ context.Context, endpointARN string) (*driver.MarketplaceEndpoint, error) { + endpoint, ok := m.marketplaceEndpoints.Get(endpointARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) + } + + result := *endpoint + + return &result, nil +} + +// ListMarketplaceModelEndpoints lists all marketplace model endpoints. +func (m *Mock) ListMarketplaceModelEndpoints(_ context.Context) ([]driver.MarketplaceEndpoint, error) { + all := m.marketplaceEndpoints.All() + out := make([]driver.MarketplaceEndpoint, 0, len(all)) + + for _, endpoint := range all { + out = append(out, *endpoint) + } + + return out, nil +} + +// UpdateMarketplaceModelEndpoint replaces an endpoint's configuration. +func (m *Mock) UpdateMarketplaceModelEndpoint( + _ context.Context, endpointARN string, endpointConfig []byte, +) (*driver.MarketplaceEndpoint, error) { + endpoint, ok := m.marketplaceEndpoints.Get(endpointARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) + } + + if len(endpointConfig) == 0 { + return nil, errors.New(errors.InvalidArgument, "endpointConfig is required") + } + + endpoint.EndpointConfig = copyBytes(endpointConfig) + endpoint.UpdatedAt = m.now() + m.marketplaceEndpoints.Set(endpointARN, endpoint) + + result := *endpoint + + return &result, nil +} + +// DeleteMarketplaceModelEndpoint deletes a marketplace model endpoint by ARN. +func (m *Mock) DeleteMarketplaceModelEndpoint(_ context.Context, endpointARN string) error { + if !m.marketplaceEndpoints.Has(endpointARN) { + return errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) + } + + m.marketplaceEndpoints.Delete(endpointARN) + + return nil +} + +// RegisterMarketplaceModelEndpoint registers an existing endpoint, marking it +// REGISTERED. +func (m *Mock) RegisterMarketplaceModelEndpoint( + _ context.Context, endpointIdentifier, modelSourceIdentifier string, +) (*driver.MarketplaceEndpoint, error) { + if modelSourceIdentifier == "" { + return nil, errors.New(errors.InvalidArgument, "modelSourceIdentifier is required") + } + + endpoint, ok := m.marketplaceEndpoints.Get(endpointIdentifier) + if !ok { + return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointIdentifier) + } + + endpoint.ModelSourceIdentifier = modelSourceIdentifier + endpoint.Status = driver.MarketplaceEndpointStatusRegistered + endpoint.UpdatedAt = m.now() + m.marketplaceEndpoints.Set(endpointIdentifier, endpoint) + + result := *endpoint + + return &result, nil +} + +// DeregisterMarketplaceModelEndpoint removes the Bedrock registration for an +// endpoint while leaving the endpoint record (and the underlying, unmodeled +// SageMaker endpoint) in place, so it can still be described or deleted. +func (m *Mock) DeregisterMarketplaceModelEndpoint(_ context.Context, endpointARN string) error { + if !m.marketplaceEndpoints.Has(endpointARN) { + return errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) + } + + return nil +} + +// --- Foundation model agreements --- + +// CreateFoundationModelAgreement records an accepted agreement for a model. +func (m *Mock) CreateFoundationModelAgreement(_ context.Context, modelID, offerToken string) (string, error) { + switch { + case modelID == "": + return "", errors.New(errors.InvalidArgument, "modelId is required") + case offerToken == "": + return "", errors.New(errors.InvalidArgument, "offerToken is required") + } + + m.fmAgreements.Set(modelID, true) + + return modelID, nil +} + +// DeleteFoundationModelAgreement removes an accepted agreement for a model. +func (m *Mock) DeleteFoundationModelAgreement(_ context.Context, modelID string) error { + if modelID == "" { + return errors.New(errors.InvalidArgument, "modelId is required") + } + + m.fmAgreements.Delete(modelID) + + return nil +} + +// ListFoundationModelAgreementOffers returns a single synthetic agreement offer. +func (*Mock) ListFoundationModelAgreementOffers( + _ context.Context, modelID, _ string, +) ([]driver.FoundationModelOffer, error) { + if modelID == "" { + return nil, errors.New(errors.InvalidArgument, "modelId is required") + } + + return []driver.FoundationModelOffer{ + {OfferToken: idgen.GenerateID(""), OfferID: idgen.GenerateID("")}, + }, nil +} + +// GetFoundationModelAvailability reports availability for a model. An accepted +// agreement makes it AUTHORIZED/AVAILABLE; otherwise NOT_AUTHORIZED/PENDING. +func (m *Mock) GetFoundationModelAvailability( + _ context.Context, modelID string, +) (*driver.FoundationModelAvailability, error) { + if modelID == "" { + return nil, errors.New(errors.InvalidArgument, "modelId is required") + } + + out := &driver.FoundationModelAvailability{ + AgreementStatus: driver.AgreementStatusPending, + AuthorizationStatus: driver.AuthorizationStatusNotAuthorized, + EntitlementAvailability: driver.AvailabilityAvailable, + RegionAvailability: driver.AvailabilityAvailable, + } + if m.fmAgreements.Has(modelID) { + out.AgreementStatus = driver.AgreementStatusAvailable + out.AuthorizationStatus = driver.AuthorizationStatusAuthorized + } + + return out, nil +} + +// copyBytes returns a copy of b so stored payloads never alias caller memory. +func copyBytes(b []byte) []byte { + if len(b) == 0 { + return nil + } + + out := make([]byte, len(b)) + copy(out, b) + + return out +} diff --git a/providers/aws/bedrock/marketplace_agreements_test.go b/providers/aws/bedrock/marketplace_agreements_test.go new file mode 100644 index 00000000..0ee0d0a1 --- /dev/null +++ b/providers/aws/bedrock/marketplace_agreements_test.go @@ -0,0 +1,139 @@ +package bedrock + +import ( + "context" + "testing" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +func TestMarketplaceEndpointLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + cfg := bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: "endpoint-1", + ModelSourceIdentifier: "arn:aws:sagemaker:us-east-1:aws:hub-content/model/1", + EndpointConfig: []byte(`{"sageMaker":{"instanceType":"ml.m5.large"}}`), + AcceptEula: true, + Tags: map[string]string{"team": "ml"}, + } + + endpoint, err := m.CreateMarketplaceModelEndpoint(ctx, cfg) + requireNoError(t, err) + assertNotEmpty(t, endpoint.EndpointARN) + assertEqual(t, bedrockdriver.MarketplaceEndpointStatusRegistered, endpoint.Status) + assertEqual(t, bedrockdriver.MarketplaceEndpointStatusInService, endpoint.EndpointStatus) + + got, err := m.GetMarketplaceModelEndpoint(ctx, endpoint.EndpointARN) + requireNoError(t, err) + assertEqual(t, cfg.ModelSourceIdentifier, got.ModelSourceIdentifier) + + tags, err := m.ListTagsForResource(ctx, endpoint.EndpointARN) + requireNoError(t, err) + assertEqual(t, 1, len(tags)) + + list, err := m.ListMarketplaceModelEndpoints(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + updated, err := m.UpdateMarketplaceModelEndpoint(ctx, endpoint.EndpointARN, []byte(`{"sageMaker":{"instanceType":"ml.m5.xlarge"}}`)) + requireNoError(t, err) + assertEqual(t, `{"sageMaker":{"instanceType":"ml.m5.xlarge"}}`, string(updated.EndpointConfig)) + + reg, err := m.RegisterMarketplaceModelEndpoint(ctx, endpoint.EndpointARN, "arn:model/source-2") + requireNoError(t, err) + assertEqual(t, bedrockdriver.MarketplaceEndpointStatusRegistered, reg.Status) + assertEqual(t, "arn:model/source-2", reg.ModelSourceIdentifier) + + requireNoError(t, m.DeregisterMarketplaceModelEndpoint(ctx, endpoint.EndpointARN)) + + requireNoError(t, m.DeleteMarketplaceModelEndpoint(ctx, endpoint.EndpointARN)) + + _, err = m.GetMarketplaceModelEndpoint(ctx, endpoint.EndpointARN) + assertError(t, err, true) +} + +func TestMarketplaceEndpointValidationAndErrors(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateMarketplaceModelEndpoint(ctx, bedrockdriver.MarketplaceEndpointConfig{ + ModelSourceIdentifier: "arn:model/s", EndpointConfig: []byte(`{}`), + }) + assertError(t, err, true) + + _, err = m.CreateMarketplaceModelEndpoint(ctx, bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: "e", EndpointConfig: []byte(`{}`), + }) + assertError(t, err, true) + + _, err = m.CreateMarketplaceModelEndpoint(ctx, bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: "e", ModelSourceIdentifier: "arn:model/s", + }) + assertError(t, err, true) + + _, err = m.GetMarketplaceModelEndpoint(ctx, "arn:missing") + assertError(t, err, true) + + _, err = m.UpdateMarketplaceModelEndpoint(ctx, "arn:missing", []byte(`{}`)) + assertError(t, err, true) + + _, err = m.RegisterMarketplaceModelEndpoint(ctx, "arn:missing", "arn:model/s") + assertError(t, err, true) + + assertError(t, m.DeregisterMarketplaceModelEndpoint(ctx, "arn:missing"), true) + assertError(t, m.DeleteMarketplaceModelEndpoint(ctx, "arn:missing"), true) +} + +func TestFoundationModelAgreementLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + avail, err := m.GetFoundationModelAvailability(ctx, titanModel) + requireNoError(t, err) + assertEqual(t, bedrockdriver.AuthorizationStatusNotAuthorized, avail.AuthorizationStatus) + assertEqual(t, bedrockdriver.AgreementStatusPending, avail.AgreementStatus) + assertEqual(t, bedrockdriver.AvailabilityAvailable, avail.EntitlementAvailability) + assertEqual(t, bedrockdriver.AvailabilityAvailable, avail.RegionAvailability) + + offers, err := m.ListFoundationModelAgreementOffers(ctx, titanModel, "ALL") + requireNoError(t, err) + assertEqual(t, 1, len(offers)) + assertNotEmpty(t, offers[0].OfferToken) + assertNotEmpty(t, offers[0].OfferID) + + modelID, err := m.CreateFoundationModelAgreement(ctx, titanModel, offers[0].OfferToken) + requireNoError(t, err) + assertEqual(t, titanModel, modelID) + + avail, err = m.GetFoundationModelAvailability(ctx, titanModel) + requireNoError(t, err) + assertEqual(t, bedrockdriver.AuthorizationStatusAuthorized, avail.AuthorizationStatus) + assertEqual(t, bedrockdriver.AgreementStatusAvailable, avail.AgreementStatus) + + requireNoError(t, m.DeleteFoundationModelAgreement(ctx, titanModel)) + + avail, err = m.GetFoundationModelAvailability(ctx, titanModel) + requireNoError(t, err) + assertEqual(t, bedrockdriver.AuthorizationStatusNotAuthorized, avail.AuthorizationStatus) +} + +func TestFoundationModelAgreementValidation(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateFoundationModelAgreement(ctx, "", "token") + assertError(t, err, true) + + _, err = m.CreateFoundationModelAgreement(ctx, titanModel, "") + assertError(t, err, true) + + _, err = m.ListFoundationModelAgreementOffers(ctx, "", "ALL") + assertError(t, err, true) + + _, err = m.GetFoundationModelAvailability(ctx, "") + assertError(t, err, true) + + assertError(t, m.DeleteFoundationModelAgreement(ctx, ""), true) +} diff --git a/providers/aws/bedrock/registries.go b/providers/aws/bedrock/registries.go new file mode 100644 index 00000000..3d4078f3 --- /dev/null +++ b/providers/aws/bedrock/registries.go @@ -0,0 +1,287 @@ +package bedrock + +import ( + "context" + "crypto/sha256" + "encoding/hex" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- Inference profiles --- + +// CreateInferenceProfile creates an application inference profile. It is ready +// immediately: recorded ACTIVE with type APPLICATION. +func (m *Mock) CreateInferenceProfile(_ context.Context, cfg driver.InferenceProfileConfig) (*driver.InferenceProfile, error) { + switch { + case cfg.Name == "": + return nil, errors.New(errors.InvalidArgument, "inferenceProfileName is required") + case cfg.ModelSourceCopyFrom == "": + return nil, errors.New(errors.InvalidArgument, "modelSource.copyFrom is required") + } + + now := m.now() + id := idgen.GenerateID("") + arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "application-inference-profile/"+id) + + profile := &driver.InferenceProfile{ + ARN: arn, + ID: id, + Name: cfg.Name, + Models: []string{cfg.ModelSourceCopyFrom}, + Status: driver.InferenceProfileStatusActive, + Type: driver.InferenceProfileTypeApplication, + Description: cfg.Description, + CreatedAt: now, + UpdatedAt: now, + } + m.inferenceProfiles.Set(id, profile) + m.setTags(arn, m.tagsFromMap(cfg.Tags)) + + result := *profile + + return &result, nil +} + +// GetInferenceProfile returns an inference profile by ID or ARN. +func (m *Mock) GetInferenceProfile(_ context.Context, identifier string) (*driver.InferenceProfile, error) { + if p, ok := m.inferenceProfiles.Get(identifier); ok { + result := *p + + return &result, nil + } + + for _, p := range m.inferenceProfiles.All() { + if p.ARN == identifier { + result := *p + + return &result, nil + } + } + + return nil, errors.Newf(errors.NotFound, "inference profile %q not found", identifier) +} + +// ListInferenceProfiles lists all inference profiles. +func (m *Mock) ListInferenceProfiles(_ context.Context) ([]driver.InferenceProfile, error) { + all := m.inferenceProfiles.All() + out := make([]driver.InferenceProfile, 0, len(all)) + + for _, p := range all { + out = append(out, *p) + } + + return out, nil +} + +// DeleteInferenceProfile deletes an inference profile by ID or ARN. +func (m *Mock) DeleteInferenceProfile(_ context.Context, identifier string) error { + if m.inferenceProfiles.Has(identifier) { + m.inferenceProfiles.Delete(identifier) + + return nil + } + + for id, p := range m.inferenceProfiles.All() { + if p.ARN == identifier { + m.inferenceProfiles.Delete(id) + + return nil + } + } + + return errors.Newf(errors.NotFound, "inference profile %q not found", identifier) +} + +// --- Prompt routers --- + +// CreatePromptRouter creates a prompt router. It is ready immediately: recorded +// AVAILABLE with type custom. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreatePromptRouter(_ context.Context, cfg driver.PromptRouterConfig) (*driver.PromptRouter, error) { + switch { + case cfg.Name == "": + return nil, errors.New(errors.InvalidArgument, "promptRouterName is required") + case len(cfg.Models) == 0: + return nil, errors.New(errors.InvalidArgument, "models is required") + case cfg.ResponseQualityDifference == nil: + return nil, errors.New(errors.InvalidArgument, "routingCriteria.responseQualityDifference is required") + case cfg.FallbackModelARN == "": + return nil, errors.New(errors.InvalidArgument, "fallbackModel.modelArn is required") + } + + now := m.now() + id := idgen.GenerateID("") + arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "prompt-router/"+id) + + router := &driver.PromptRouter{ + ARN: arn, + Name: cfg.Name, + Models: append([]string(nil), cfg.Models...), + ResponseQualityDifference: cfg.ResponseQualityDifference, + FallbackModelARN: cfg.FallbackModelARN, + Status: driver.PromptRouterStatusAvailable, + Type: driver.PromptRouterTypeCustom, + Description: cfg.Description, + CreatedAt: now, + UpdatedAt: now, + } + m.promptRouters.Set(arn, router) + m.setTags(arn, m.tagsFromMap(cfg.Tags)) + + result := *router + + return &result, nil +} + +// GetPromptRouter returns a prompt router by its ARN. +func (m *Mock) GetPromptRouter(_ context.Context, promptRouterARN string) (*driver.PromptRouter, error) { + router, ok := m.promptRouters.Get(promptRouterARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "prompt router %q not found", promptRouterARN) + } + + result := *router + + return &result, nil +} + +// ListPromptRouters lists all prompt routers. +func (m *Mock) ListPromptRouters(_ context.Context) ([]driver.PromptRouter, error) { + all := m.promptRouters.All() + out := make([]driver.PromptRouter, 0, len(all)) + + for _, router := range all { + out = append(out, *router) + } + + return out, nil +} + +// DeletePromptRouter deletes a prompt router by its ARN. +func (m *Mock) DeletePromptRouter(_ context.Context, promptRouterARN string) error { + if !m.promptRouters.Has(promptRouterARN) { + return errors.Newf(errors.NotFound, "prompt router %q not found", promptRouterARN) + } + + m.promptRouters.Delete(promptRouterARN) + + return nil +} + +// --- Automated reasoning policies --- + +// CreateAutomatedReasoningPolicy creates an automated reasoning policy at the +// DRAFT version. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateAutomatedReasoningPolicy( + _ context.Context, cfg driver.AutomatedReasoningPolicyConfig, +) (*driver.AutomatedReasoningPolicy, error) { + if cfg.Name == "" { + return nil, errors.New(errors.InvalidArgument, "name is required") + } + + now := m.now() + id := idgen.GenerateID("") + arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "automated-reasoning-policy/"+id) + + policy := &driver.AutomatedReasoningPolicy{ + ARN: arn, + ID: id, + Name: cfg.Name, + Version: driver.AutomatedReasoningPolicyVersionDraft, + DefinitionHash: definitionHash(cfg.PolicyDefinition), + Description: cfg.Description, + PolicyDefinition: cfg.PolicyDefinition, + CreatedAt: now, + UpdatedAt: now, + } + if cfg.KMSKeyID != "" { + policy.KMSKeyARN = cfg.KMSKeyID + } + + m.arPolicies.Set(arn, policy) + m.setTags(arn, m.tagsFromMap(cfg.Tags)) + + result := *policy + + return &result, nil +} + +// GetAutomatedReasoningPolicy returns an automated reasoning policy by its ARN. +func (m *Mock) GetAutomatedReasoningPolicy(_ context.Context, policyARN string) (*driver.AutomatedReasoningPolicy, error) { + policy, ok := m.arPolicies.Get(policyARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "automated reasoning policy %q not found", policyARN) + } + + result := *policy + + return &result, nil +} + +// ListAutomatedReasoningPolicies lists all automated reasoning policies. +func (m *Mock) ListAutomatedReasoningPolicies(_ context.Context) ([]driver.AutomatedReasoningPolicy, error) { + all := m.arPolicies.All() + out := make([]driver.AutomatedReasoningPolicy, 0, len(all)) + + for _, policy := range all { + out = append(out, *policy) + } + + return out, nil +} + +// UpdateAutomatedReasoningPolicy updates an existing policy's definition, name, +// and description, refreshing its definition hash and update time. +func (m *Mock) UpdateAutomatedReasoningPolicy( + _ context.Context, policyARN string, upd driver.AutomatedReasoningPolicyUpdate, +) (*driver.AutomatedReasoningPolicy, error) { + policy, ok := m.arPolicies.Get(policyARN) + if !ok { + return nil, errors.Newf(errors.NotFound, "automated reasoning policy %q not found", policyARN) + } + + if upd.Name != "" { + policy.Name = upd.Name + } + + if upd.Description != "" { + policy.Description = upd.Description + } + + if len(upd.PolicyDefinition) != 0 { + policy.PolicyDefinition = upd.PolicyDefinition + } + + policy.DefinitionHash = definitionHash(policy.PolicyDefinition) + policy.UpdatedAt = m.now() + m.arPolicies.Set(policyARN, policy) + + result := *policy + + return &result, nil +} + +// DeleteAutomatedReasoningPolicy deletes an automated reasoning policy by its ARN. +func (m *Mock) DeleteAutomatedReasoningPolicy(_ context.Context, policyARN string) error { + if !m.arPolicies.Has(policyARN) { + return errors.Newf(errors.NotFound, "automated reasoning policy %q not found", policyARN) + } + + m.arPolicies.Delete(policyARN) + + return nil +} + +// definitionHash returns a deterministic non-empty hash of a policy definition, +// used as the policy's concurrency token. +func definitionHash(def []byte) string { + sum := sha256.Sum256(def) + + return hex.EncodeToString(sum[:]) +} diff --git a/providers/aws/bedrock/registries_test.go b/providers/aws/bedrock/registries_test.go new file mode 100644 index 00000000..9ecb1095 --- /dev/null +++ b/providers/aws/bedrock/registries_test.go @@ -0,0 +1,181 @@ +package bedrock + +import ( + "context" + "testing" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +func TestInferenceProfileLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + src := "arn:aws:bedrock:us-east-1:123456789012:foundation-model/" + titanModel + profile, err := m.CreateInferenceProfile(ctx, bedrockdriver.InferenceProfileConfig{ + Name: "profile-1", + ModelSourceCopyFrom: src, + Description: "test profile", + Tags: map[string]string{"team": "ml"}, + }) + requireNoError(t, err) + assertNotEmpty(t, profile.ARN) + assertNotEmpty(t, profile.ID) + assertEqual(t, bedrockdriver.InferenceProfileStatusActive, profile.Status) + assertEqual(t, bedrockdriver.InferenceProfileTypeApplication, profile.Type) + assertEqual(t, 1, len(profile.Models)) + assertEqual(t, src, profile.Models[0]) + + byID, err := m.GetInferenceProfile(ctx, profile.ID) + requireNoError(t, err) + assertEqual(t, profile.ARN, byID.ARN) + + byARN, err := m.GetInferenceProfile(ctx, profile.ARN) + requireNoError(t, err) + assertEqual(t, "profile-1", byARN.Name) + + list, err := m.ListInferenceProfiles(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + tags, err := m.ListTagsForResource(ctx, profile.ARN) + requireNoError(t, err) + assertEqual(t, 1, len(tags)) + + requireNoError(t, m.DeleteInferenceProfile(ctx, profile.ID)) + + _, err = m.GetInferenceProfile(ctx, profile.ID) + assertError(t, err, true) +} + +func TestInferenceProfileValidation(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateInferenceProfile(ctx, bedrockdriver.InferenceProfileConfig{ModelSourceCopyFrom: "arn:model"}) + assertError(t, err, true) + + _, err = m.CreateInferenceProfile(ctx, bedrockdriver.InferenceProfileConfig{Name: "p"}) + assertError(t, err, true) + + assertError(t, m.DeleteInferenceProfile(ctx, "missing"), true) +} + +func TestPromptRouterLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + diff := 0.5 + router, err := m.CreatePromptRouter(ctx, bedrockdriver.PromptRouterConfig{ + Name: "router-1", + Models: []string{"arn:aws:bedrock:us-east-1::foundation-model/a", "arn:model/b"}, + ResponseQualityDifference: &diff, + FallbackModelARN: "arn:model/fallback", + Description: "test router", + Tags: map[string]string{"team": "ml"}, + }) + requireNoError(t, err) + assertNotEmpty(t, router.ARN) + assertEqual(t, bedrockdriver.PromptRouterStatusAvailable, router.Status) + assertEqual(t, bedrockdriver.PromptRouterTypeCustom, router.Type) + assertEqual(t, 2, len(router.Models)) + + got, err := m.GetPromptRouter(ctx, router.ARN) + requireNoError(t, err) + assertEqual(t, "router-1", got.Name) + assertEqual(t, "arn:model/fallback", got.FallbackModelARN) + + if got.ResponseQualityDifference == nil || *got.ResponseQualityDifference != diff { + t.Fatalf("unexpected routing criteria: %+v", got.ResponseQualityDifference) + } + + list, err := m.ListPromptRouters(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + requireNoError(t, m.DeletePromptRouter(ctx, router.ARN)) + + _, err = m.GetPromptRouter(ctx, router.ARN) + assertError(t, err, true) +} + +func TestPromptRouterValidation(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + diff := 0.25 + _, err := m.CreatePromptRouter(ctx, bedrockdriver.PromptRouterConfig{ + Models: []string{"arn:model/a"}, + ResponseQualityDifference: &diff, + FallbackModelARN: "arn:model/f", + }) + assertError(t, err, true) + + _, err = m.CreatePromptRouter(ctx, bedrockdriver.PromptRouterConfig{Name: "r", FallbackModelARN: "arn:model/f", ResponseQualityDifference: &diff}) + assertError(t, err, true) + + _, err = m.CreatePromptRouter(ctx, bedrockdriver.PromptRouterConfig{Name: "r", Models: []string{"arn:model/a"}, FallbackModelARN: "arn:model/f"}) + assertError(t, err, true) + + assertError(t, m.DeletePromptRouter(ctx, "arn:missing"), true) +} + +func TestAutomatedReasoningPolicyLifecycle(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + policy, err := m.CreateAutomatedReasoningPolicy(ctx, bedrockdriver.AutomatedReasoningPolicyConfig{ + Name: "policy-1", + Description: "test policy", + KMSKeyID: "arn:aws:kms:us-east-1:123456789012:key/abc", + PolicyDefinition: []byte(`{"rules":[]}`), + Tags: map[string]string{"team": "ml"}, + }) + requireNoError(t, err) + assertNotEmpty(t, policy.ARN) + assertNotEmpty(t, policy.ID) + assertNotEmpty(t, policy.DefinitionHash) + assertEqual(t, bedrockdriver.AutomatedReasoningPolicyVersionDraft, policy.Version) + assertEqual(t, "arn:aws:kms:us-east-1:123456789012:key/abc", policy.KMSKeyARN) + + got, err := m.GetAutomatedReasoningPolicy(ctx, policy.ARN) + requireNoError(t, err) + assertEqual(t, "policy-1", got.Name) + + list, err := m.ListAutomatedReasoningPolicies(ctx) + requireNoError(t, err) + assertEqual(t, 1, len(list)) + + updated, err := m.UpdateAutomatedReasoningPolicy(ctx, policy.ARN, bedrockdriver.AutomatedReasoningPolicyUpdate{ + Name: "policy-1-renamed", + Description: "updated", + PolicyDefinition: []byte(`{"rules":[{"id":"r1"}]}`), + }) + requireNoError(t, err) + assertEqual(t, "policy-1-renamed", updated.Name) + + if updated.DefinitionHash == policy.DefinitionHash { + t.Fatal("expected definition hash to change after update") + } + + requireNoError(t, m.DeleteAutomatedReasoningPolicy(ctx, policy.ARN)) + + _, err = m.GetAutomatedReasoningPolicy(ctx, policy.ARN) + assertError(t, err, true) +} + +func TestAutomatedReasoningPolicyValidationAndErrors(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateAutomatedReasoningPolicy(ctx, bedrockdriver.AutomatedReasoningPolicyConfig{Description: "no name"}) + assertError(t, err, true) + + _, err = m.GetAutomatedReasoningPolicy(ctx, "arn:missing") + assertError(t, err, true) + + _, err = m.UpdateAutomatedReasoningPolicy(ctx, "arn:missing", bedrockdriver.AutomatedReasoningPolicyUpdate{Name: "x"}) + assertError(t, err, true) + + assertError(t, m.DeleteAutomatedReasoningPolicy(ctx, "arn:missing"), true) +} diff --git a/providers/aws/bedrock/tags.go b/providers/aws/bedrock/tags.go new file mode 100644 index 00000000..c20f2fc2 --- /dev/null +++ b/providers/aws/bedrock/tags.go @@ -0,0 +1,110 @@ +package bedrock + +import ( + "context" + "sort" + + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// TagResource merges tags onto the resource ARN's tag list, overwriting values +// for existing keys. Tagging an ARN the emulator does not track is a no-op +// success, mirroring AWS. +func (m *Mock) TagResource(_ context.Context, resourceARN string, tags []driver.Tag) error { + existing, _ := m.tags.Get(resourceARN) + + merged := make([]driver.Tag, 0, len(existing)+len(tags)) + index := map[string]int{} + + for _, t := range existing { + index[t.Key] = len(merged) + merged = append(merged, t) + } + + for _, t := range tags { + if i, ok := index[t.Key]; ok { + merged[i].Value = t.Value + + continue + } + + index[t.Key] = len(merged) + merged = append(merged, t) + } + + m.tags.Set(resourceARN, merged) + + return nil +} + +// UntagResource removes the given tag keys from the resource ARN's tag list. +func (m *Mock) UntagResource(_ context.Context, resourceARN string, tagKeys []string) error { + existing, ok := m.tags.Get(resourceARN) + if !ok { + return nil + } + + remove := make(map[string]bool, len(tagKeys)) + for _, k := range tagKeys { + remove[k] = true + } + + kept := make([]driver.Tag, 0, len(existing)) + + for _, t := range existing { + if !remove[t.Key] { + kept = append(kept, t) + } + } + + m.tags.Set(resourceARN, kept) + + return nil +} + +// ListTagsForResource returns a defensive copy of the resource ARN's tags in +// stable insertion order (an empty slice when none are set). +func (m *Mock) ListTagsForResource(_ context.Context, resourceARN string) ([]driver.Tag, error) { + tags, _ := m.tags.Get(resourceARN) + + return copyTags(tags), nil +} + +// setTags records the initial tags for a resource ARN (no-op for empty input). +func (m *Mock) setTags(arn string, tags []driver.Tag) { + if len(tags) == 0 { + return + } + + m.tags.Set(arn, copyTags(tags)) +} + +// tagsFromMap converts a string map into a []driver.Tag sorted by key, giving +// deterministic ordering for create-time seeding. +func (*Mock) tagsFromMap(tags map[string]string) []driver.Tag { + if len(tags) == 0 { + return nil + } + + keys := make([]string, 0, len(tags)) + for k := range tags { + keys = append(keys, k) + } + + sort.Strings(keys) + + out := make([]driver.Tag, 0, len(keys)) + for _, k := range keys { + out = append(out, driver.Tag{Key: k, Value: tags[k]}) + } + + return out +} + +// copyTags returns a defensive copy of tags (nil-safe, always non-nil). +func copyTags(tags []driver.Tag) []driver.Tag { + out := make([]driver.Tag, len(tags)) + copy(out, tags) + + return out +} diff --git a/providers/aws/bedrock/tags_test.go b/providers/aws/bedrock/tags_test.go new file mode 100644 index 00000000..635aa0e9 --- /dev/null +++ b/providers/aws/bedrock/tags_test.go @@ -0,0 +1,96 @@ +package bedrock + +import ( + "context" + "testing" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +const tagResourceARN = "arn:aws:bedrock:us-east-1:123456789012:guardrail/gr-test" + +func TestTagResourceAndList(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + err := m.TagResource(ctx, tagResourceARN, []bedrockdriver.Tag{ + {Key: "env", Value: "prod"}, + {Key: "team", Value: "ml"}, + }) + requireNoError(t, err) + + tags, err := m.ListTagsForResource(ctx, tagResourceARN) + requireNoError(t, err) + assertEqual(t, 2, len(tags)) + assertEqual(t, "env", tags[0].Key) + assertEqual(t, "prod", tags[0].Value) + assertEqual(t, "team", tags[1].Key) + assertEqual(t, "ml", tags[1].Value) +} + +func TestTagResourceOverwrite(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + requireNoError(t, m.TagResource(ctx, tagResourceARN, []bedrockdriver.Tag{{Key: "env", Value: "dev"}})) + requireNoError(t, m.TagResource(ctx, tagResourceARN, []bedrockdriver.Tag{{Key: "env", Value: "prod"}})) + + tags, err := m.ListTagsForResource(ctx, tagResourceARN) + requireNoError(t, err) + assertEqual(t, 1, len(tags)) + assertEqual(t, "env", tags[0].Key) + assertEqual(t, "prod", tags[0].Value) +} + +func TestUntagResource(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + requireNoError(t, m.TagResource(ctx, tagResourceARN, []bedrockdriver.Tag{ + {Key: "env", Value: "prod"}, + {Key: "team", Value: "ml"}, + {Key: "cost", Value: "42"}, + })) + + requireNoError(t, m.UntagResource(ctx, tagResourceARN, []string{"team", "cost"})) + + tags, err := m.ListTagsForResource(ctx, tagResourceARN) + requireNoError(t, err) + assertEqual(t, 1, len(tags)) + assertEqual(t, "env", tags[0].Key) +} + +func TestListTagsUnknownARN(t *testing.T) { + m := newTestMock() + + tags, err := m.ListTagsForResource(context.Background(), "arn:aws:bedrock:us-east-1:123456789012:guardrail/nope") + requireNoError(t, err) + + if tags == nil { + t.Fatal("expected a non-nil empty slice for an unknown ARN") + } + + assertEqual(t, 0, len(tags)) +} + +func TestCreateGuardrailPersistsTags(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + g, err := m.CreateGuardrail(ctx, bedrockdriver.GuardrailConfig{ + Name: "gr-tagged", + BlockedInputMessaging: "blocked in", + BlockedOutputsMessaging: "blocked out", + Tags: map[string]string{"env": "prod", "team": "ml"}, + }) + requireNoError(t, err) + + tags, err := m.ListTagsForResource(ctx, g.ARN) + requireNoError(t, err) + assertEqual(t, 2, len(tags)) + // tagsFromMap sorts by key, so ordering is deterministic. + assertEqual(t, "env", tags[0].Key) + assertEqual(t, "prod", tags[0].Value) + assertEqual(t, "team", tags[1].Key) + assertEqual(t, "ml", tags[1].Value) +} diff --git a/providers/aws/bedrockagent/agents.go b/providers/aws/bedrockagent/agents.go new file mode 100644 index 00000000..153f580f --- /dev/null +++ b/providers/aws/bedrockagent/agents.go @@ -0,0 +1,168 @@ +package bedrockagent + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// CreateAgent creates an agent in the NOT_PREPARED state. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateAgent(_ context.Context, cfg driver.AgentConfig) (*driver.Agent, error) { + if cfg.Name == "" { + return nil, errors.New(errors.InvalidArgument, "agentName is required") + } + + id := idgen.GenerateID("AGENT") + now := m.now() + + ttl := cfg.IdleSessionTTLInSeconds + if ttl == 0 { + ttl = defaultIdleSessionTTL + } + + agent := &driver.Agent{ + ID: id, + ARN: idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "agent/"+id), + Name: cfg.Name, + ResourceRoleArn: cfg.ResourceRoleArn, + FoundationModel: cfg.FoundationModel, + Instruction: cfg.Instruction, + Description: cfg.Description, + Status: driver.AgentNotPrepared, + Version: driver.DraftVersion, + IdleSessionTTLInSeconds: ttl, + CreatedAt: now, + UpdatedAt: now, + } + m.agents.Set(id, agent) + + result := *agent + + return &result, nil +} + +// GetAgent returns an agent by ID. +func (m *Mock) GetAgent(_ context.Context, agentID string) (*driver.Agent, error) { + agent, ok := m.agents.Get(agentID) + if !ok { + return nil, errors.Newf(errors.NotFound, "agent %q not found", agentID) + } + + result := *agent + + return &result, nil +} + +// ListAgents lists all agents. +func (m *Mock) ListAgents(_ context.Context) ([]driver.Agent, error) { + all := m.agents.SortedValues() + out := make([]driver.Agent, 0, len(all)) + + for _, a := range all { + out = append(out, *a) + } + + return out, nil +} + +// UpdateAgent updates an agent's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) UpdateAgent(_ context.Context, agentID string, cfg driver.AgentConfig) (*driver.Agent, error) { + agent, ok := m.agents.Get(agentID) + if !ok { + return nil, errors.Newf(errors.NotFound, "agent %q not found", agentID) + } + + updated := *agent + updated.Name = orDefault(cfg.Name, agent.Name) + updated.ResourceRoleArn = orDefault(cfg.ResourceRoleArn, agent.ResourceRoleArn) + updated.FoundationModel = orDefault(cfg.FoundationModel, agent.FoundationModel) + updated.Instruction = orDefault(cfg.Instruction, agent.Instruction) + updated.Description = cfg.Description + updated.Status = driver.AgentNotPrepared + updated.UpdatedAt = m.now() + + if cfg.IdleSessionTTLInSeconds != 0 { + updated.IdleSessionTTLInSeconds = cfg.IdleSessionTTLInSeconds + } + + m.agents.Set(agentID, &updated) + + result := updated + + return &result, nil +} + +// DeleteAgent deletes an agent and returns its terminal status. +func (m *Mock) DeleteAgent(_ context.Context, agentID string) (string, error) { + if !m.agents.Has(agentID) { + return "", errors.Newf(errors.NotFound, "agent %q not found", agentID) + } + + m.agents.Delete(agentID) + + return statusDeleting, nil +} + +// PrepareAgent prepares an agent, transitioning it to PREPARED. +func (m *Mock) PrepareAgent(_ context.Context, agentID string) (*driver.Agent, error) { + agent, ok := m.agents.Get(agentID) + if !ok { + return nil, errors.Newf(errors.NotFound, "agent %q not found", agentID) + } + + updated := *agent + updated.Status = driver.AgentPrepared + updated.PreparedAt = m.now() + updated.UpdatedAt = updated.PreparedAt + m.agents.Set(agentID, &updated) + + result := updated + + return &result, nil +} + +// CreateAgentAlias creates an alias of an agent in the PREPARED state. +func (m *Mock) CreateAgentAlias(_ context.Context, cfg driver.AgentAliasConfig) (*driver.AgentAlias, error) { + switch { + case cfg.AgentID == "": + return nil, errors.New(errors.InvalidArgument, "agentId is required") + case cfg.Name == "": + return nil, errors.New(errors.InvalidArgument, "agentAliasName is required") + } + + if !m.agents.Has(cfg.AgentID) { + return nil, errors.Newf(errors.NotFound, "agent %q not found", cfg.AgentID) + } + + id := idgen.GenerateID("ALIAS") + now := m.now() + alias := &driver.AgentAlias{ + ID: id, + ARN: idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "agent-alias/"+cfg.AgentID+"/"+id), + AgentID: cfg.AgentID, + Name: cfg.Name, + Description: cfg.Description, + Status: driver.AgentAliasPrepared, + CreatedAt: now, + UpdatedAt: now, + } + m.aliases.Set(id, alias) + + result := *alias + + return &result, nil +} + +func orDefault(v, fallback string) string { + if v == "" { + return fallback + } + + return v +} diff --git a/providers/aws/bedrockagent/bedrockagent.go b/providers/aws/bedrockagent/bedrockagent.go new file mode 100644 index 00000000..df3439e7 --- /dev/null +++ b/providers/aws/bedrockagent/bedrockagent.go @@ -0,0 +1,53 @@ +// Package bedrockagent provides an in-memory mock implementation of the AWS +// Bedrock Agent authoring control plane: agents (and aliases), knowledge bases, +// data sources, ingestion jobs, flows, and prompts. Resources are created +// directly in a terminal/ready state. +package bedrockagent + +import ( + "time" + + "github.com/stackshy/cloudemu/v2/config" + "github.com/stackshy/cloudemu/v2/internal/memstore" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// defaultIdleSessionTTL is the idle-session timeout assigned to agents created +// without an explicit value, matching the Bedrock default. +const defaultIdleSessionTTL int32 = 600 + +// statusDeleting is the transitional status returned by delete operations. +const statusDeleting = "DELETING" + +// Compile-time check that Mock implements driver.BedrockAgent. +var _ driver.BedrockAgent = (*Mock)(nil) + +// Mock is an in-memory mock implementation of the AWS Bedrock Agent service. +type Mock struct { + agents *memstore.Store[*driver.Agent] + aliases *memstore.Store[*driver.AgentAlias] + knowledge *memstore.Store[*driver.KnowledgeBase] + dataSource *memstore.Store[*driver.DataSource] + jobs *memstore.Store[*driver.IngestionJob] + flows *memstore.Store[*driver.Flow] + prompts *memstore.Store[*driver.Prompt] + opts *config.Options +} + +// New creates a new Bedrock Agent mock. +func New(opts *config.Options) *Mock { + return &Mock{ + agents: memstore.New[*driver.Agent](), + aliases: memstore.New[*driver.AgentAlias](), + knowledge: memstore.New[*driver.KnowledgeBase](), + dataSource: memstore.New[*driver.DataSource](), + jobs: memstore.New[*driver.IngestionJob](), + flows: memstore.New[*driver.Flow](), + prompts: memstore.New[*driver.Prompt](), + opts: opts, + } +} + +func (m *Mock) now() string { + return m.opts.Clock.Now().UTC().Format(time.RFC3339) +} diff --git a/providers/aws/bedrockagent/bedrockagent_test.go b/providers/aws/bedrockagent/bedrockagent_test.go new file mode 100644 index 00000000..3607ca02 --- /dev/null +++ b/providers/aws/bedrockagent/bedrockagent_test.go @@ -0,0 +1,155 @@ +package bedrockagent + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/stackshy/cloudemu/v2/config" + cerrors "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +func newMock() *Mock { + clock := config.NewFakeClock(time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)) + + return New(config.NewOptions(config.WithClock(clock))) +} + +func TestAgentLifecycle(t *testing.T) { + m := newMock() + ctx := context.Background() + + agent, err := m.CreateAgent(ctx, driver.AgentConfig{Name: "a1", FoundationModel: "fm"}) + require.NoError(t, err) + assert.Equal(t, driver.AgentNotPrepared, agent.Status) + assert.Equal(t, driver.DraftVersion, agent.Version) + assert.Equal(t, defaultIdleSessionTTL, agent.IdleSessionTTLInSeconds) + assert.NotEmpty(t, agent.ARN) + + prepared, err := m.PrepareAgent(ctx, agent.ID) + require.NoError(t, err) + assert.Equal(t, driver.AgentPrepared, prepared.Status) + assert.NotEmpty(t, prepared.PreparedAt) + + updated, err := m.UpdateAgent(ctx, agent.ID, driver.AgentConfig{Name: "a1-new"}) + require.NoError(t, err) + assert.Equal(t, "a1-new", updated.Name) + assert.Equal(t, "fm", updated.FoundationModel) // preserved + + agents, err := m.ListAgents(ctx) + require.NoError(t, err) + assert.Len(t, agents, 1) + + _, err = m.DeleteAgent(ctx, agent.ID) + require.NoError(t, err) + + _, err = m.GetAgent(ctx, agent.ID) + assert.True(t, cerrors.IsNotFound(err)) +} + +func TestCreateAgentValidation(t *testing.T) { + _, err := newMock().CreateAgent(context.Background(), driver.AgentConfig{}) + assert.True(t, cerrors.IsInvalidArgument(err)) +} + +func TestAgentAliasRequiresAgent(t *testing.T) { + m := newMock() + + _, err := m.CreateAgentAlias(context.Background(), driver.AgentAliasConfig{AgentID: "nope", Name: "x"}) + assert.True(t, cerrors.IsNotFound(err)) +} + +func TestKnowledgeBaseAndDataSource(t *testing.T) { + m := newMock() + ctx := context.Background() + + kb, err := m.CreateKnowledgeBase(ctx, driver.KnowledgeBaseConfig{ + Name: "kb1", + RoleArn: "role", + KnowledgeBaseConfiguration: json.RawMessage(`{"type":"VECTOR"}`), + }) + require.NoError(t, err) + assert.Equal(t, driver.KnowledgeBaseActive, kb.Status) + + ds, err := m.CreateDataSource(ctx, driver.DataSourceConfig{ + KnowledgeBaseID: kb.ID, + Name: "ds1", + DataSourceConfiguration: json.RawMessage(`{"type":"S3"}`), + }) + require.NoError(t, err) + assert.Equal(t, driver.DataSourceAvailable, ds.Status) + assert.Equal(t, kb.ID, ds.KnowledgeBaseID) + + job, err := m.StartIngestionJob(ctx, kb.ID, ds.ID, "reindex") + require.NoError(t, err) + assert.Equal(t, driver.IngestionJobComplete, job.Status) + + sources, err := m.ListDataSources(ctx, kb.ID) + require.NoError(t, err) + assert.Len(t, sources, 1) + + _, err = m.DeleteDataSource(ctx, kb.ID, ds.ID) + require.NoError(t, err) + + _, err = m.GetDataSource(ctx, kb.ID, ds.ID) + assert.True(t, cerrors.IsNotFound(err)) +} + +func TestDataSourceRequiresKnowledgeBase(t *testing.T) { + m := newMock() + + _, err := m.CreateDataSource(context.Background(), driver.DataSourceConfig{ + KnowledgeBaseID: "missing", + Name: "ds1", + DataSourceConfiguration: json.RawMessage(`{"type":"S3"}`), + }) + assert.True(t, cerrors.IsNotFound(err)) +} + +func TestFlowLifecycle(t *testing.T) { + m := newMock() + ctx := context.Background() + + flow, err := m.CreateFlow(ctx, driver.FlowConfig{Name: "f1", ExecutionRoleArn: "role"}) + require.NoError(t, err) + assert.Equal(t, driver.FlowNotPrepared, flow.Status) + + prepared, err := m.PrepareFlow(ctx, flow.ID) + require.NoError(t, err) + assert.Equal(t, driver.FlowPrepared, prepared.Status) + + id, err := m.DeleteFlow(ctx, flow.ID) + require.NoError(t, err) + assert.Equal(t, flow.ID, id) + + _, err = m.GetFlow(ctx, flow.ID) + assert.True(t, cerrors.IsNotFound(err)) +} + +func TestPromptLifecycle(t *testing.T) { + m := newMock() + ctx := context.Background() + + prompt, err := m.CreatePrompt(ctx, driver.PromptConfig{Name: "p1"}) + require.NoError(t, err) + assert.Equal(t, driver.DraftVersion, prompt.Version) + + updated, err := m.UpdatePrompt(ctx, prompt.ID, driver.PromptConfig{Name: "p1-new"}) + require.NoError(t, err) + assert.Equal(t, "p1-new", updated.Name) + + prompts, err := m.ListPrompts(ctx) + require.NoError(t, err) + assert.Len(t, prompts, 1) + + _, err = m.DeletePrompt(ctx, prompt.ID) + require.NoError(t, err) + + _, err = m.GetPrompt(ctx, prompt.ID) + assert.True(t, cerrors.IsNotFound(err)) +} diff --git a/providers/aws/bedrockagent/datasources.go b/providers/aws/bedrockagent/datasources.go new file mode 100644 index 00000000..cd83f0bf --- /dev/null +++ b/providers/aws/bedrockagent/datasources.go @@ -0,0 +1,142 @@ +package bedrockagent + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// CreateDataSource creates a data source under a knowledge base in the +// AVAILABLE state. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateDataSource(_ context.Context, cfg driver.DataSourceConfig) (*driver.DataSource, error) { + switch { + case cfg.Name == "": + return nil, errors.New(errors.InvalidArgument, "name is required") + case len(cfg.DataSourceConfiguration) == 0: + return nil, errors.New(errors.InvalidArgument, "dataSourceConfiguration is required") + } + + if !m.knowledge.Has(cfg.KnowledgeBaseID) { + return nil, errors.Newf(errors.NotFound, "knowledge base %q not found", cfg.KnowledgeBaseID) + } + + id := idgen.GenerateID("DS") + now := m.now() + ds := &driver.DataSource{ + ID: id, + KnowledgeBaseID: cfg.KnowledgeBaseID, + Name: cfg.Name, + Description: cfg.Description, + Status: driver.DataSourceAvailable, + DataDeletionPolicy: cfg.DataDeletionPolicy, + DataSourceConfiguration: cfg.DataSourceConfiguration, + CreatedAt: now, + UpdatedAt: now, + } + m.dataSource.Set(id, ds) + + result := *ds + + return &result, nil +} + +// GetDataSource returns a data source by knowledge-base and data-source ID. +func (m *Mock) GetDataSource(_ context.Context, kbID, dsID string) (*driver.DataSource, error) { + ds := m.findDataSource(kbID, dsID) + if ds == nil { + return nil, errors.Newf(errors.NotFound, "data source %q not found", dsID) + } + + result := *ds + + return &result, nil +} + +// ListDataSources lists all data sources under a knowledge base. +func (m *Mock) ListDataSources(_ context.Context, kbID string) ([]driver.DataSource, error) { + all := m.dataSource.SortedValues() + out := make([]driver.DataSource, 0, len(all)) + + for _, ds := range all { + if ds.KnowledgeBaseID == kbID { + out = append(out, *ds) + } + } + + return out, nil +} + +// UpdateDataSource updates a data source's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) UpdateDataSource(_ context.Context, cfg driver.DataSourceConfig, dsID string) (*driver.DataSource, error) { + ds := m.findDataSource(cfg.KnowledgeBaseID, dsID) + if ds == nil { + return nil, errors.Newf(errors.NotFound, "data source %q not found", dsID) + } + + updated := *ds + updated.Name = orDefault(cfg.Name, ds.Name) + updated.Description = cfg.Description + updated.DataDeletionPolicy = orDefault(cfg.DataDeletionPolicy, ds.DataDeletionPolicy) + updated.UpdatedAt = m.now() + + if len(cfg.DataSourceConfiguration) != 0 { + updated.DataSourceConfiguration = cfg.DataSourceConfiguration + } + + m.dataSource.Set(dsID, &updated) + + result := updated + + return &result, nil +} + +// DeleteDataSource deletes a data source and returns its terminal status. +func (m *Mock) DeleteDataSource(_ context.Context, kbID, dsID string) (string, error) { + if m.findDataSource(kbID, dsID) == nil { + return "", errors.Newf(errors.NotFound, "data source %q not found", dsID) + } + + m.dataSource.Delete(dsID) + + return statusDeleting, nil +} + +// StartIngestionJob starts an ingestion job that completes synchronously. +func (m *Mock) StartIngestionJob(_ context.Context, kbID, dsID, description string) (*driver.IngestionJob, error) { + if m.findDataSource(kbID, dsID) == nil { + return nil, errors.Newf(errors.NotFound, "data source %q not found", dsID) + } + + id := idgen.GenerateID("JOB") + now := m.now() + job := &driver.IngestionJob{ + ID: id, + KnowledgeBaseID: kbID, + DataSourceID: dsID, + Description: description, + Status: driver.IngestionJobComplete, + StartedAt: now, + UpdatedAt: now, + } + m.jobs.Set(id, job) + + result := *job + + return &result, nil +} + +// findDataSource returns the data source matching dsID scoped to kbID, or nil. +func (m *Mock) findDataSource(kbID, dsID string) *driver.DataSource { + ds, ok := m.dataSource.Get(dsID) + if !ok || ds.KnowledgeBaseID != kbID { + return nil + } + + return ds +} diff --git a/providers/aws/bedrockagent/flows.go b/providers/aws/bedrockagent/flows.go new file mode 100644 index 00000000..37a6f5d0 --- /dev/null +++ b/providers/aws/bedrockagent/flows.go @@ -0,0 +1,121 @@ +package bedrockagent + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// CreateFlow creates a flow in the NotPrepared state. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateFlow(_ context.Context, cfg driver.FlowConfig) (*driver.Flow, error) { + switch { + case cfg.Name == "": + return nil, errors.New(errors.InvalidArgument, "name is required") + case cfg.ExecutionRoleArn == "": + return nil, errors.New(errors.InvalidArgument, "executionRoleArn is required") + } + + id := idgen.GenerateID("FLOW") + now := m.now() + flow := &driver.Flow{ + ID: id, + ARN: idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "flow/"+id), + Name: cfg.Name, + ExecutionRoleArn: cfg.ExecutionRoleArn, + Description: cfg.Description, + Status: driver.FlowNotPrepared, + Version: driver.DraftVersion, + CustomerEncryptionKeyArn: cfg.CustomerEncryptionKeyArn, + Definition: cfg.Definition, + CreatedAt: now, + UpdatedAt: now, + } + m.flows.Set(id, flow) + + result := *flow + + return &result, nil +} + +// GetFlow returns a flow by identifier. +func (m *Mock) GetFlow(_ context.Context, id string) (*driver.Flow, error) { + flow, ok := m.flows.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "flow %q not found", id) + } + + result := *flow + + return &result, nil +} + +// ListFlows lists all flows. +func (m *Mock) ListFlows(_ context.Context) ([]driver.Flow, error) { + all := m.flows.SortedValues() + out := make([]driver.Flow, 0, len(all)) + + for _, f := range all { + out = append(out, *f) + } + + return out, nil +} + +// UpdateFlow updates a flow's mutable fields, resetting it to NotPrepared. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) UpdateFlow(_ context.Context, id string, cfg driver.FlowConfig) (*driver.Flow, error) { + flow, ok := m.flows.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "flow %q not found", id) + } + + updated := *flow + updated.Name = orDefault(cfg.Name, flow.Name) + updated.ExecutionRoleArn = orDefault(cfg.ExecutionRoleArn, flow.ExecutionRoleArn) + updated.Description = cfg.Description + updated.Status = driver.FlowNotPrepared + updated.UpdatedAt = m.now() + + if len(cfg.Definition) != 0 { + updated.Definition = cfg.Definition + } + + m.flows.Set(id, &updated) + + result := updated + + return &result, nil +} + +// DeleteFlow deletes a flow and returns its identifier. +func (m *Mock) DeleteFlow(_ context.Context, id string) (string, error) { + if !m.flows.Has(id) { + return "", errors.Newf(errors.NotFound, "flow %q not found", id) + } + + m.flows.Delete(id) + + return id, nil +} + +// PrepareFlow prepares a flow, transitioning it to Prepared. +func (m *Mock) PrepareFlow(_ context.Context, id string) (*driver.Flow, error) { + flow, ok := m.flows.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "flow %q not found", id) + } + + updated := *flow + updated.Status = driver.FlowPrepared + updated.UpdatedAt = m.now() + m.flows.Set(id, &updated) + + result := updated + + return &result, nil +} diff --git a/providers/aws/bedrockagent/knowledgebases.go b/providers/aws/bedrockagent/knowledgebases.go new file mode 100644 index 00000000..7f085994 --- /dev/null +++ b/providers/aws/bedrockagent/knowledgebases.go @@ -0,0 +1,108 @@ +package bedrockagent + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// CreateKnowledgeBase creates a knowledge base in the ACTIVE state. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreateKnowledgeBase(_ context.Context, cfg driver.KnowledgeBaseConfig) (*driver.KnowledgeBase, error) { + switch { + case cfg.Name == "": + return nil, errors.New(errors.InvalidArgument, "name is required") + case cfg.RoleArn == "": + return nil, errors.New(errors.InvalidArgument, "roleArn is required") + case len(cfg.KnowledgeBaseConfiguration) == 0: + return nil, errors.New(errors.InvalidArgument, "knowledgeBaseConfiguration is required") + } + + id := idgen.GenerateID("KB") + now := m.now() + kb := &driver.KnowledgeBase{ + ID: id, + ARN: idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "knowledge-base/"+id), + Name: cfg.Name, + RoleArn: cfg.RoleArn, + Description: cfg.Description, + Status: driver.KnowledgeBaseActive, + KnowledgeBaseConfiguration: cfg.KnowledgeBaseConfiguration, + StorageConfiguration: cfg.StorageConfiguration, + CreatedAt: now, + UpdatedAt: now, + } + m.knowledge.Set(id, kb) + + result := *kb + + return &result, nil +} + +// GetKnowledgeBase returns a knowledge base by ID. +func (m *Mock) GetKnowledgeBase(_ context.Context, id string) (*driver.KnowledgeBase, error) { + kb, ok := m.knowledge.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "knowledge base %q not found", id) + } + + result := *kb + + return &result, nil +} + +// ListKnowledgeBases lists all knowledge bases. +func (m *Mock) ListKnowledgeBases(_ context.Context) ([]driver.KnowledgeBase, error) { + all := m.knowledge.SortedValues() + out := make([]driver.KnowledgeBase, 0, len(all)) + + for _, kb := range all { + out = append(out, *kb) + } + + return out, nil +} + +// UpdateKnowledgeBase updates a knowledge base's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) UpdateKnowledgeBase(_ context.Context, id string, cfg driver.KnowledgeBaseConfig) (*driver.KnowledgeBase, error) { + kb, ok := m.knowledge.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "knowledge base %q not found", id) + } + + updated := *kb + updated.Name = orDefault(cfg.Name, kb.Name) + updated.RoleArn = orDefault(cfg.RoleArn, kb.RoleArn) + updated.Description = cfg.Description + updated.UpdatedAt = m.now() + + if len(cfg.KnowledgeBaseConfiguration) != 0 { + updated.KnowledgeBaseConfiguration = cfg.KnowledgeBaseConfiguration + } + + if len(cfg.StorageConfiguration) != 0 { + updated.StorageConfiguration = cfg.StorageConfiguration + } + + m.knowledge.Set(id, &updated) + + result := updated + + return &result, nil +} + +// DeleteKnowledgeBase deletes a knowledge base and returns its terminal status. +func (m *Mock) DeleteKnowledgeBase(_ context.Context, id string) (string, error) { + if !m.knowledge.Has(id) { + return "", errors.Newf(errors.NotFound, "knowledge base %q not found", id) + } + + m.knowledge.Delete(id) + + return statusDeleting, nil +} diff --git a/providers/aws/bedrockagent/prompts.go b/providers/aws/bedrockagent/prompts.go new file mode 100644 index 00000000..5a29e07f --- /dev/null +++ b/providers/aws/bedrockagent/prompts.go @@ -0,0 +1,99 @@ +package bedrockagent + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// CreatePrompt creates a prompt at the DRAFT version. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) CreatePrompt(_ context.Context, cfg driver.PromptConfig) (*driver.Prompt, error) { + if cfg.Name == "" { + return nil, errors.New(errors.InvalidArgument, "name is required") + } + + id := idgen.GenerateID("PROMPT") + now := m.now() + prompt := &driver.Prompt{ + ID: id, + ARN: idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "prompt/"+id), + Name: cfg.Name, + Description: cfg.Description, + Version: driver.DraftVersion, + DefaultVariant: cfg.DefaultVariant, + CustomerEncryptionKeyArn: cfg.CustomerEncryptionKeyArn, + Variants: cfg.Variants, + CreatedAt: now, + UpdatedAt: now, + } + m.prompts.Set(id, prompt) + + result := *prompt + + return &result, nil +} + +// GetPrompt returns a prompt by identifier. +func (m *Mock) GetPrompt(_ context.Context, id string) (*driver.Prompt, error) { + prompt, ok := m.prompts.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "prompt %q not found", id) + } + + result := *prompt + + return &result, nil +} + +// ListPrompts lists all prompts. +func (m *Mock) ListPrompts(_ context.Context) ([]driver.Prompt, error) { + all := m.prompts.SortedValues() + out := make([]driver.Prompt, 0, len(all)) + + for _, p := range all { + out = append(out, *p) + } + + return out, nil +} + +// UpdatePrompt updates a prompt's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (m *Mock) UpdatePrompt(_ context.Context, id string, cfg driver.PromptConfig) (*driver.Prompt, error) { + prompt, ok := m.prompts.Get(id) + if !ok { + return nil, errors.Newf(errors.NotFound, "prompt %q not found", id) + } + + updated := *prompt + updated.Name = orDefault(cfg.Name, prompt.Name) + updated.Description = cfg.Description + updated.DefaultVariant = orDefault(cfg.DefaultVariant, prompt.DefaultVariant) + updated.UpdatedAt = m.now() + + if len(cfg.Variants) != 0 { + updated.Variants = cfg.Variants + } + + m.prompts.Set(id, &updated) + + result := updated + + return &result, nil +} + +// DeletePrompt deletes a prompt and returns its identifier. +func (m *Mock) DeletePrompt(_ context.Context, id string) (string, error) { + if !m.prompts.Has(id) { + return "", errors.Newf(errors.NotFound, "prompt %q not found", id) + } + + m.prompts.Delete(id) + + return id, nil +} diff --git a/providers/aws/bedrockagentruntime/bedrockagentruntime.go b/providers/aws/bedrockagentruntime/bedrockagentruntime.go new file mode 100644 index 00000000..3c0cf5ac --- /dev/null +++ b/providers/aws/bedrockagentruntime/bedrockagentruntime.go @@ -0,0 +1,109 @@ +// Package bedrockagentruntime provides an in-memory mock implementation of the +// AWS Bedrock Agent runtime (bedrock-agent-runtime): InvokeAgent, Retrieve, and +// RetrieveAndGenerate. Responses are deterministic so real SDK callers get +// stable, parseable output. +package bedrockagentruntime + +import ( + "context" + "fmt" + + "github.com/stackshy/cloudemu/v2/config" + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/internal/idgen" + "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +// Compile-time check that Mock implements driver.BedrockAgentRuntime. +var _ driver.BedrockAgentRuntime = (*Mock)(nil) + +const contentTypeJSON = "application/json" + +// Mock is an in-memory mock implementation of the AWS Bedrock Agent runtime. +type Mock struct { + opts *config.Options +} + +// New creates a new Bedrock Agent runtime mock. +func New(opts *config.Options) *Mock { + return &Mock{opts: opts} +} + +// InvokeAgent returns a deterministic simulated agent completion. The response +// echoes the session id from the request path; when the caller omits it a fresh +// session id is generated. +func (*Mock) InvokeAgent(_ context.Context, in driver.InvokeAgentInput) (*driver.InvokeAgentResult, error) { + if in.AgentID == "" { + return nil, errors.New(errors.InvalidArgument, "agentId is required") + } + + if in.AgentAliasID == "" { + return nil, errors.New(errors.InvalidArgument, "agentAliasId is required") + } + + sessionID := in.SessionID + if sessionID == "" { + sessionID = idgen.GenerateID("session-") + } + + var completion string + if in.InputText == "" { + completion = fmt.Sprintf("This is a simulated agent response from agent %s.", in.AgentID) + } else { + completion = fmt.Sprintf("This is a simulated agent response to: %s", in.InputText) + } + + return &driver.InvokeAgentResult{ + Completion: completion, + SessionID: sessionID, + ContentType: contentTypeJSON, + }, nil +} + +// Retrieve returns deterministic fake chunks that echo the query text. +func (m *Mock) Retrieve(_ context.Context, in driver.RetrieveInput) (*driver.RetrieveResult, error) { + if in.KnowledgeBaseID == "" { + return nil, errors.New(errors.InvalidArgument, "knowledgeBaseId is required") + } + + if in.QueryText == "" { + return nil, errors.New(errors.InvalidArgument, "retrievalQuery.text is required") + } + + base := fmt.Sprintf("s3://cloudemu-%s/kb/%s", m.opts.Region, in.KnowledgeBaseID) + + return &driver.RetrieveResult{ + Results: []driver.RetrievalResult{ + { + Text: fmt.Sprintf("Simulated knowledge-base result 1 for query: %s", in.QueryText), + Score: 0.95, + LocationURI: base + "/doc-1.txt", + }, + { + Text: fmt.Sprintf("Simulated knowledge-base result 2 for query: %s", in.QueryText), + Score: 0.82, + LocationURI: base + "/doc-2.txt", + }, + }, + }, nil +} + +// RetrieveAndGenerate returns a deterministic generated answer and a session id. +// A supplied session id is reused; otherwise a new one is generated. +func (*Mock) RetrieveAndGenerate( + _ context.Context, in driver.RetrieveAndGenerateInput, +) (*driver.RetrieveAndGenerateResult, error) { + if in.InputText == "" { + return nil, errors.New(errors.InvalidArgument, "input.text is required") + } + + sessionID := in.SessionID + if sessionID == "" { + sessionID = idgen.GenerateID("rag-session-") + } + + return &driver.RetrieveAndGenerateResult{ + Text: fmt.Sprintf("This is a simulated retrieve-and-generate answer to: %s", in.InputText), + SessionID: sessionID, + }, nil +} diff --git a/providers/aws/bedrockagentruntime/bedrockagentruntime_test.go b/providers/aws/bedrockagentruntime/bedrockagentruntime_test.go new file mode 100644 index 00000000..f2ba7ceb --- /dev/null +++ b/providers/aws/bedrockagentruntime/bedrockagentruntime_test.go @@ -0,0 +1,158 @@ +package bedrockagentruntime + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/stackshy/cloudemu/v2/config" + "github.com/stackshy/cloudemu/v2/errors" + "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +func newTestMock() *Mock { + fc := config.NewFakeClock(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)) + opts := config.NewOptions( + config.WithClock(fc), + config.WithRegion("us-east-1"), + config.WithAccountID("123456789012"), + ) + + return New(opts) +} + +func TestInvokeAgent(t *testing.T) { + m := newTestMock() + + out, err := m.InvokeAgent(context.Background(), driver.InvokeAgentInput{ + AgentID: "AGENT123", + AgentAliasID: "ALIAS123", + SessionID: "sess-1", + InputText: "hello agent", + }) + if err != nil { + t.Fatalf("InvokeAgent: %v", err) + } + + if !strings.Contains(out.Completion, "hello agent") { + t.Fatalf("completion %q does not echo the prompt", out.Completion) + } + + if out.SessionID != "sess-1" { + t.Fatalf("got session id %q, want sess-1", out.SessionID) + } + + if out.ContentType != "application/json" { + t.Fatalf("got content type %q, want application/json", out.ContentType) + } +} + +func TestInvokeAgentGeneratesSessionID(t *testing.T) { + m := newTestMock() + + out, err := m.InvokeAgent(context.Background(), driver.InvokeAgentInput{ + AgentID: "AGENT123", + AgentAliasID: "ALIAS123", + }) + if err != nil { + t.Fatalf("InvokeAgent: %v", err) + } + + if out.SessionID == "" { + t.Fatal("expected a generated session id") + } +} + +func TestInvokeAgentValidation(t *testing.T) { + m := newTestMock() + + if _, err := m.InvokeAgent(context.Background(), driver.InvokeAgentInput{AgentAliasID: "a"}); !errors.IsInvalidArgument(err) { + t.Fatalf("want InvalidArgument for missing agentId, got %v", err) + } + + if _, err := m.InvokeAgent(context.Background(), driver.InvokeAgentInput{AgentID: "a"}); !errors.IsInvalidArgument(err) { + t.Fatalf("want InvalidArgument for missing agentAliasId, got %v", err) + } +} + +func TestRetrieve(t *testing.T) { + m := newTestMock() + + out, err := m.Retrieve(context.Background(), driver.RetrieveInput{ + KnowledgeBaseID: "KB123", + QueryText: "what is bedrock", + }) + if err != nil { + t.Fatalf("Retrieve: %v", err) + } + + if len(out.Results) == 0 { + t.Fatal("expected non-empty retrieval results") + } + + for _, r := range out.Results { + if !strings.Contains(r.Text, "what is bedrock") { + t.Fatalf("result %q does not echo the query", r.Text) + } + + if r.LocationURI == "" { + t.Fatal("expected a location uri") + } + } +} + +func TestRetrieveValidation(t *testing.T) { + m := newTestMock() + + if _, err := m.Retrieve(context.Background(), driver.RetrieveInput{KnowledgeBaseID: "KB"}); !errors.IsInvalidArgument(err) { + t.Fatalf("want InvalidArgument for missing query text, got %v", err) + } + + if _, err := m.Retrieve(context.Background(), driver.RetrieveInput{QueryText: "x"}); !errors.IsInvalidArgument(err) { + t.Fatalf("want InvalidArgument for missing knowledgeBaseId, got %v", err) + } +} + +func TestRetrieveAndGenerate(t *testing.T) { + m := newTestMock() + + out, err := m.RetrieveAndGenerate(context.Background(), driver.RetrieveAndGenerateInput{ + InputText: "summarize bedrock", + }) + if err != nil { + t.Fatalf("RetrieveAndGenerate: %v", err) + } + + if !strings.Contains(out.Text, "summarize bedrock") { + t.Fatalf("generated text %q does not echo the input", out.Text) + } + + if out.SessionID == "" { + t.Fatal("expected a generated session id") + } +} + +func TestRetrieveAndGenerateReusesSessionID(t *testing.T) { + m := newTestMock() + + out, err := m.RetrieveAndGenerate(context.Background(), driver.RetrieveAndGenerateInput{ + InputText: "hi", + SessionID: "existing-session", + }) + if err != nil { + t.Fatalf("RetrieveAndGenerate: %v", err) + } + + if out.SessionID != "existing-session" { + t.Fatalf("got session id %q, want existing-session", out.SessionID) + } +} + +func TestRetrieveAndGenerateValidation(t *testing.T) { + m := newTestMock() + + if _, err := m.RetrieveAndGenerate(context.Background(), driver.RetrieveAndGenerateInput{}); !errors.IsInvalidArgument(err) { + t.Fatalf("want InvalidArgument for missing input text, got %v", err) + } +} diff --git a/server/aws/aws.go b/server/aws/aws.go index e71c668c..34a7c752 100644 --- a/server/aws/aws.go +++ b/server/aws/aws.go @@ -11,6 +11,8 @@ import ( eksdriver "github.com/stackshy/cloudemu/v2/providers/aws/eks/driver" "github.com/stackshy/cloudemu/v2/server" "github.com/stackshy/cloudemu/v2/server/aws/bedrock" + "github.com/stackshy/cloudemu/v2/server/aws/bedrockagent" + "github.com/stackshy/cloudemu/v2/server/aws/bedrockagentruntime" "github.com/stackshy/cloudemu/v2/server/aws/cloudwatch" cloudwatchlogssrv "github.com/stackshy/cloudemu/v2/server/aws/cloudwatchlogs" "github.com/stackshy/cloudemu/v2/server/aws/dynamodb" @@ -35,6 +37,8 @@ import ( ssmsrv "github.com/stackshy/cloudemu/v2/server/aws/ssm" stssrv "github.com/stackshy/cloudemu/v2/server/aws/sts" bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" + bedrockagentdriver "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" + bedrockagentruntimedriver "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" cachedriver "github.com/stackshy/cloudemu/v2/services/cache/driver" computedriver "github.com/stackshy/cloudemu/v2/services/compute/driver" crdriver "github.com/stackshy/cloudemu/v2/services/containerregistry/driver" @@ -62,20 +66,22 @@ import ( // field nil to omit that service; the server returns 501 Not Implemented for // any request that no registered handler matches. type Drivers struct { - S3 storagedriver.Bucket - DynamoDB dbdriver.Database - EC2 computedriver.Compute - VPC netdriver.Networking - CloudWatch mondriver.Monitoring - Lambda sdrv.Serverless - SQS mqdriver.MessageQueue - RDS rdbdriver.RelationalDB - Redshift rdbdriver.RelationalDB - EKS eksdriver.EKS - IAM iamdriver.IAM - ECR crdriver.ContainerRegistry - Bedrock bedrockdriver.Bedrock - SageMaker sagemakerdriver.Service + S3 storagedriver.Bucket + DynamoDB dbdriver.Database + EC2 computedriver.Compute + VPC netdriver.Networking + CloudWatch mondriver.Monitoring + Lambda sdrv.Serverless + SQS mqdriver.MessageQueue + RDS rdbdriver.RelationalDB + Redshift rdbdriver.RelationalDB + EKS eksdriver.EKS + IAM iamdriver.IAM + ECR crdriver.ContainerRegistry + Bedrock bedrockdriver.Bedrock + BedrockAgent bedrockagentdriver.BedrockAgent + BedrockAgentRuntime bedrockagentruntimedriver.BedrockAgentRuntime + SageMaker sagemakerdriver.Service // SecretsManager serves the Secrets Manager JSON 1.1 protocol against // the secrets driver. SecretsManager secretsdriver.Secrets @@ -124,33 +130,35 @@ type Drivers struct { // nil for the caller to inject when a shared cluster is desired. func DriversFrom(p *awsprovider.Provider) Drivers { return Drivers{ - S3: p.S3, - DynamoDB: p.DynamoDB, - EC2: p.EC2, - VPC: p.VPC, - CloudWatch: p.CloudWatch, - Lambda: p.Lambda, - SQS: p.SQS, - RDS: p.RDS, - Redshift: p.Redshift, - EKS: p.EKS, - IAM: p.IAM, - ECR: p.ECR, - Bedrock: p.Bedrock, - SageMaker: p.SageMaker, - SecretsManager: p.SecretsManager, - SSM: p.SSM, - CloudWatchLogs: p.CloudWatchLogs, - Route53: p.Route53, - ELB: p.ELB, - EventBridge: p.EventBridge, - ElastiCache: p.ElastiCache, - SNS: p.SNS, - STS: true, - K8sAPI: nil, // injected by the caller when a shared cluster is desired - ResourceDiscovery: p.ResourceDiscovery, - AccountID: p.AccountID, - Region: p.Region, + S3: p.S3, + DynamoDB: p.DynamoDB, + EC2: p.EC2, + VPC: p.VPC, + CloudWatch: p.CloudWatch, + Lambda: p.Lambda, + SQS: p.SQS, + RDS: p.RDS, + Redshift: p.Redshift, + EKS: p.EKS, + IAM: p.IAM, + ECR: p.ECR, + Bedrock: p.Bedrock, + BedrockAgent: p.BedrockAgent, + BedrockAgentRuntime: p.BedrockAgentRuntime, + SageMaker: p.SageMaker, + SecretsManager: p.SecretsManager, + SSM: p.SSM, + CloudWatchLogs: p.CloudWatchLogs, + Route53: p.Route53, + ELB: p.ELB, + EventBridge: p.EventBridge, + ElastiCache: p.ElastiCache, + SNS: p.SNS, + STS: true, + K8sAPI: nil, // injected by the caller when a shared cluster is desired + ResourceDiscovery: p.ResourceDiscovery, + AccountID: p.AccountID, + Region: p.Region, } } @@ -312,6 +320,22 @@ func New(d Drivers) *server.Server { srv.Register(bedrock.New(d.Bedrock)) } + // bedrock-agent-runtime (InvokeAgent / Retrieve / RetrieveAndGenerate) shares + // the /agents and /knowledgebases roots with the bedrock-agent control plane, + // but matches only the runtime suffixes (/text, /retrieve) and + // /retrieveAndGenerate. It MUST register before the control-plane handler so + // its more specific Matches wins for those paths, and both before S3. + if d.BedrockAgentRuntime != nil { + srv.Register(bedrockagentruntime.New(d.BedrockAgentRuntime)) + } + + // bedrock-agent control plane: agents, knowledge bases, data sources, flows, + // prompts. REST/JSON rooted at /agents, /knowledgebases, /flows, /prompts — + // registered before S3's permissive REST fallback. + if d.BedrockAgent != nil { + srv.Register(bedrockagent.New(d.BedrockAgent)) + } + // SageMaker control plane matches the X-Amz-Target prefix "SageMaker." // (disjoint from DynamoDB/SQS/Resource-Groups-Tagging), and the runtime // matches /endpoints/{name}/invocations. The runtime path must register diff --git a/server/aws/bedrock/asyncinvoke.go b/server/aws/bedrock/asyncinvoke.go new file mode 100644 index 00000000..feb7d460 --- /dev/null +++ b/server/aws/bedrock/asyncinvoke.go @@ -0,0 +1,210 @@ +package bedrock + +import ( + "encoding/json" + "net/http" + "strings" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- wire types --- + +type asyncS3OutputConfig struct { + S3URI string `json:"s3Uri"` + BucketOwner string `json:"bucketOwner,omitempty"` + KMSKeyID string `json:"kmsKeyId,omitempty"` +} + +type asyncOutputDataConfig struct { + S3OutputDataConfig *asyncS3OutputConfig `json:"s3OutputDataConfig,omitempty"` +} + +type startAsyncInvokeRequest struct { + ClientRequestToken string `json:"clientRequestToken"` + ModelID string `json:"modelId"` + ModelInput json.RawMessage `json:"modelInput"` + OutputDataConfig *asyncOutputDataConfig `json:"outputDataConfig"` + Tags []tagPair `json:"tags"` +} + +type startAsyncInvokeResponse struct { + InvocationARN string `json:"invocationArn"` +} + +type asyncInvokeJSON struct { + InvocationARN string `json:"invocationArn"` + ModelARN string `json:"modelArn"` + Status string `json:"status"` + SubmitTime string `json:"submitTime,omitempty"` + OutputDataConfig *asyncOutputDataConfig `json:"outputDataConfig,omitempty"` + ClientRequestToken string `json:"clientRequestToken,omitempty"` + EndTime string `json:"endTime,omitempty"` + LastModifiedTime string `json:"lastModifiedTime,omitempty"` + FailureMessage string `json:"failureMessage,omitempty"` +} + +// asyncInvokeSummaryJSON mirrors the SDK's AsyncInvokeSummary shape. +type asyncInvokeSummaryJSON struct { + InvocationARN string `json:"invocationArn"` + ModelARN string `json:"modelArn"` + Status string `json:"status"` + SubmitTime string `json:"submitTime,omitempty"` + OutputDataConfig *asyncOutputDataConfig `json:"outputDataConfig,omitempty"` + ClientRequestToken string `json:"clientRequestToken,omitempty"` + EndTime string `json:"endTime,omitempty"` + LastModifiedTime string `json:"lastModifiedTime,omitempty"` + FailureMessage string `json:"failureMessage,omitempty"` +} + +type listAsyncInvokesResponse struct { + AsyncInvokeSummaries []asyncInvokeSummaryJSON `json:"asyncInvokeSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- dispatchers --- + +// serveAsyncJobs routes the async-invoke and control-plane job surfaces. Split +// out of ServeHTTP to keep each dispatcher small. +func (h *Handler) serveAsyncJobs(w http.ResponseWriter, r *http.Request, p string) { + switch { + case p == prefixAsyncInvoke || strings.HasPrefix(p, prefixAsyncInvoke+"/"): + h.serveAsyncInvoke(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixAsyncInvoke), "/")) + case p == prefixImportJobs || strings.HasPrefix(p, prefixImportJobs+"/"): + h.serveImportJobs(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixImportJobs), "/")) + case p == prefixCopyJobs || strings.HasPrefix(p, prefixCopyJobs+"/"): + h.serveCopyJobs(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixCopyJobs), "/")) + case p == prefixEvalJobs || strings.HasPrefix(p, prefixEvalJobs+"/"): + h.serveEvalJobs(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixEvalJobs), "/")) + case strings.HasPrefix(p, prefixEvalJobStop): + h.serveEvalJobStop(w, r, strings.TrimPrefix(p, prefixEvalJobStop)) + default: + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported path: "+p) + } +} + +// serveAsyncInvoke handles /async-invoke[/{invocationArn}]. The invocation ARN +// contains slashes, so it is the entire remainder of the path. +func (h *Handler) serveAsyncInvoke(w http.ResponseWriter, r *http.Request, arn string) { + if arn == "" { + switch r.Method { + case http.MethodPost: + h.startAsyncInvoke(w, r) + case http.MethodGet: + h.listAsyncInvokes(w, r) + default: + methodNotAllowed(w) + } + + return + } + + if r.Method != http.MethodGet { + methodNotAllowed(w) + + return + } + + h.getAsyncInvoke(w, r, arn) +} + +// --- operations --- + +func (h *Handler) startAsyncInvoke(w http.ResponseWriter, r *http.Request) { + var in startAsyncInvokeRequest + if !decodeJSON(w, r, &in) { + return + } + + inv, err := h.bedrock.StartAsyncInvoke(r.Context(), bedrockdriver.StartAsyncInvokeConfig{ + ClientRequestToken: in.ClientRequestToken, + ModelID: in.ModelID, + ModelInput: []byte(in.ModelInput), + Output: toDriverAsyncOutput(in.OutputDataConfig), + Tags: tagsToMap(in.Tags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, startAsyncInvokeResponse{InvocationARN: inv.InvocationARN}) +} + +func (h *Handler) getAsyncInvoke(w http.ResponseWriter, r *http.Request, arn string) { + inv, err := h.bedrock.GetAsyncInvoke(r.Context(), arn) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toAsyncInvokeJSON(inv)) +} + +func (h *Handler) listAsyncInvokes(w http.ResponseWriter, r *http.Request) { + invs, err := h.bedrock.ListAsyncInvokes(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]asyncInvokeSummaryJSON, 0, len(invs)) + for i := range invs { + out = append(out, toAsyncInvokeSummaryJSON(&invs[i])) + } + + writeJSON(w, listAsyncInvokesResponse{AsyncInvokeSummaries: out}) +} + +// --- converters --- + +func toDriverAsyncOutput(in *asyncOutputDataConfig) bedrockdriver.AsyncInvokeOutputConfig { + if in == nil || in.S3OutputDataConfig == nil { + return bedrockdriver.AsyncInvokeOutputConfig{} + } + + s3 := in.S3OutputDataConfig + + return bedrockdriver.AsyncInvokeOutputConfig{S3URI: s3.S3URI, BucketOwner: s3.BucketOwner, KMSKeyID: s3.KMSKeyID} +} + +func toAsyncOutputJSON(o bedrockdriver.AsyncInvokeOutputConfig) *asyncOutputDataConfig { + if o.S3URI == "" { + return nil + } + + return &asyncOutputDataConfig{ + S3OutputDataConfig: &asyncS3OutputConfig{S3URI: o.S3URI, BucketOwner: o.BucketOwner, KMSKeyID: o.KMSKeyID}, + } +} + +func toAsyncInvokeJSON(inv *bedrockdriver.AsyncInvoke) asyncInvokeJSON { + return asyncInvokeJSON{ + InvocationARN: inv.InvocationARN, + ModelARN: inv.ModelARN, + Status: inv.Status, + SubmitTime: inv.SubmitTime, + OutputDataConfig: toAsyncOutputJSON(inv.Output), + ClientRequestToken: inv.ClientRequestToken, + EndTime: inv.EndTime, + LastModifiedTime: inv.LastModifiedTime, + FailureMessage: inv.FailureMessage, + } +} + +func toAsyncInvokeSummaryJSON(inv *bedrockdriver.AsyncInvoke) asyncInvokeSummaryJSON { + return asyncInvokeSummaryJSON{ + InvocationARN: inv.InvocationARN, + ModelARN: inv.ModelARN, + Status: inv.Status, + SubmitTime: inv.SubmitTime, + OutputDataConfig: toAsyncOutputJSON(inv.Output), + ClientRequestToken: inv.ClientRequestToken, + EndTime: inv.EndTime, + LastModifiedTime: inv.LastModifiedTime, + FailureMessage: inv.FailureMessage, + } +} diff --git a/server/aws/bedrock/counttokens_applyguardrail.go b/server/aws/bedrock/counttokens_applyguardrail.go new file mode 100644 index 00000000..41eab21b --- /dev/null +++ b/server/aws/bedrock/counttokens_applyguardrail.go @@ -0,0 +1,94 @@ +package bedrock + +import ( + "net/http" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// countTokensBody wraps the countTokensRequest union under the "input" key the +// real bedrockruntime SDK emits ({"input":{"converse"|"invokeModel":...}}). +type countTokensBody struct { + Input countTokensRequest `json:"input"` +} + +// countTokens handles POST /model/{modelId}/count-tokens. The request body is a +// union carrying either a converse or an invokeModel member. +func (h *Handler) countTokens(w http.ResponseWriter, r *http.Request, modelID string) { + var body countTokensBody + if !decodeJSON(w, r, &body) { + return + } + + req := body.Input + in := bedrockdriver.CountTokensInput{ModelID: modelID} + + switch { + case req.InvokeModel != nil: + in.InvokeBody = req.InvokeModel.Body + case req.Converse != nil: + in.Messages = toDriverMessages(req.Converse.Messages) + in.System = textsOf(req.Converse.System) + } + + n, err := h.bedrock.CountTokens(r.Context(), in) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, countTokensResponse{InputTokens: n}) +} + +// applyGuardrail handles POST +// /guardrail/{guardrailIdentifier}/version/{guardrailVersion}/apply. The +// identifier and version arrive as path parameters. +func (h *Handler) applyGuardrail(w http.ResponseWriter, r *http.Request, guardrailID, guardrailVersion string) { + var req applyGuardrailRequest + if !decodeJSON(w, r, &req) { + return + } + + contents := make([]string, 0, len(req.Content)) + + for _, block := range req.Content { + if block.Text != nil { + contents = append(contents, block.Text.Text) + } + } + + out, err := h.bedrock.ApplyGuardrail(r.Context(), bedrockdriver.ApplyGuardrailInput{ + GuardrailIdentifier: guardrailID, + GuardrailVersion: guardrailVersion, + Source: req.Source, + Content: contents, + }) + if err != nil { + writeErr(w, err) + + return + } + + outputs := make([]applyGuardrailOutputContent, 0, len(out.Outputs)) + for _, o := range out.Outputs { + outputs = append(outputs, applyGuardrailOutputContent{Text: o}) + } + + writeJSON(w, applyGuardrailResponse{ + Usage: applyGuardrailUsage{}, + Action: out.Action, + Outputs: outputs, + Assessments: []applyGuardrailAssessment{}, + }) +} + +// toDriverMessages maps converse wire messages to driver messages. +func toDriverMessages(msgs []converseMessage) []bedrockdriver.Message { + out := make([]bedrockdriver.Message, 0, len(msgs)) + for _, m := range msgs { + out = append(out, bedrockdriver.Message{Role: m.Role, Text: textsOf(m.Content)}) + } + + return out +} diff --git a/server/aws/bedrock/guardrail_policies.go b/server/aws/bedrock/guardrail_policies.go new file mode 100644 index 00000000..354c3789 --- /dev/null +++ b/server/aws/bedrock/guardrail_policies.go @@ -0,0 +1,284 @@ +package bedrock + +import bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" + +// --- leaf wire types (shared between request and response: identical keys) --- + +type topicJSON struct { + Name string `json:"name"` + Definition string `json:"definition"` + Examples []string `json:"examples,omitempty"` + Type string `json:"type,omitempty"` +} + +type contentFilterJSON struct { + Type string `json:"type"` + InputStrength string `json:"inputStrength"` + OutputStrength string `json:"outputStrength"` +} + +type wordJSON struct { + Text string `json:"text"` +} + +type managedWordsJSON struct { + Type string `json:"type"` +} + +type piiEntityJSON struct { + Type string `json:"type"` + Action string `json:"action"` +} + +type regexJSON struct { + Name string `json:"name"` + Pattern string `json:"pattern"` + Action string `json:"action"` + Description string `json:"description,omitempty"` +} + +type contextualGroundingFilterJSON struct { + Type string `json:"type"` + Threshold float64 `json:"threshold"` + Action string `json:"action,omitempty"` +} + +// --- request wrappers ("...Config"-suffixed keys) --- + +type topicPolicyConfigJSON struct { + TopicsConfig []topicJSON `json:"topicsConfig"` +} + +type contentPolicyConfigJSON struct { + FiltersConfig []contentFilterJSON `json:"filtersConfig"` +} + +type wordPolicyConfigJSON struct { + WordsConfig []wordJSON `json:"wordsConfig,omitempty"` + ManagedWordListsConfig []managedWordsJSON `json:"managedWordListsConfig,omitempty"` +} + +type sensitiveInfoPolicyConfigJSON struct { + PiiEntitiesConfig []piiEntityJSON `json:"piiEntitiesConfig,omitempty"` + RegexesConfig []regexJSON `json:"regexesConfig,omitempty"` +} + +type contextualGroundingPolicyConfigJSON struct { + FiltersConfig []contextualGroundingFilterJSON `json:"filtersConfig"` +} + +// --- response wrappers ("Config" suffix dropped) --- + +type topicPolicyJSON struct { + Topics []topicJSON `json:"topics"` +} + +type contentPolicyJSON struct { + Filters []contentFilterJSON `json:"filters"` +} + +type wordPolicyJSON struct { + Words []wordJSON `json:"words,omitempty"` + ManagedWordLists []managedWordsJSON `json:"managedWordLists,omitempty"` +} + +type sensitiveInfoPolicyJSON struct { + PiiEntities []piiEntityJSON `json:"piiEntities,omitempty"` + Regexes []regexJSON `json:"regexes,omitempty"` +} + +type contextualGroundingPolicyJSON struct { + Filters []contextualGroundingFilterJSON `json:"filters"` +} + +// --- request -> driver conversion --- + +func toDriverGuardrailPolicies(in *createGuardrailRequest) bedrockdriver.GuardrailPolicies { + var p bedrockdriver.GuardrailPolicies + + if c := in.TopicPolicyConfig; c != nil { + p.TopicPolicy = &bedrockdriver.GuardrailTopicPolicy{Topics: toDriverTopics(c.TopicsConfig)} + } + + if c := in.ContentPolicyConfig; c != nil { + p.ContentPolicy = &bedrockdriver.GuardrailContentPolicy{Filters: toDriverContentFilters(c.FiltersConfig)} + } + + if c := in.WordPolicyConfig; c != nil { + p.WordPolicy = &bedrockdriver.GuardrailWordPolicy{ + Words: toDriverWords(c.WordsConfig), + ManagedWordLists: toDriverManagedWords(c.ManagedWordListsConfig), + } + } + + if c := in.SensitiveInformationPolicyConfig; c != nil { + p.SensitiveInformationPolicy = &bedrockdriver.GuardrailSensitiveInformationPolicy{ + PiiEntities: toDriverPiiEntities(c.PiiEntitiesConfig), + Regexes: toDriverRegexes(c.RegexesConfig), + } + } + + if c := in.ContextualGroundingPolicyConfig; c != nil { + p.ContextualGroundingPolicy = &bedrockdriver.GuardrailContextualGroundingPolicy{ + Filters: toDriverGroundingFilters(c.FiltersConfig), + } + } + + return p +} + +func toDriverTopics(in []topicJSON) []bedrockdriver.GuardrailTopic { + out := make([]bedrockdriver.GuardrailTopic, len(in)) + for i, t := range in { + out[i] = bedrockdriver.GuardrailTopic{Name: t.Name, Definition: t.Definition, Examples: t.Examples, Type: t.Type} + } + + return out +} + +func toDriverContentFilters(in []contentFilterJSON) []bedrockdriver.GuardrailContentFilter { + out := make([]bedrockdriver.GuardrailContentFilter, len(in)) + for i, f := range in { + out[i] = bedrockdriver.GuardrailContentFilter{Type: f.Type, InputStrength: f.InputStrength, OutputStrength: f.OutputStrength} + } + + return out +} + +func toDriverWords(in []wordJSON) []bedrockdriver.GuardrailWord { + out := make([]bedrockdriver.GuardrailWord, len(in)) + for i, w := range in { + out[i] = bedrockdriver.GuardrailWord{Text: w.Text} + } + + return out +} + +func toDriverManagedWords(in []managedWordsJSON) []bedrockdriver.GuardrailManagedWordList { + out := make([]bedrockdriver.GuardrailManagedWordList, len(in)) + for i, w := range in { + out[i] = bedrockdriver.GuardrailManagedWordList{Type: w.Type} + } + + return out +} + +func toDriverPiiEntities(in []piiEntityJSON) []bedrockdriver.GuardrailPiiEntity { + out := make([]bedrockdriver.GuardrailPiiEntity, len(in)) + for i, e := range in { + out[i] = bedrockdriver.GuardrailPiiEntity{Type: e.Type, Action: e.Action} + } + + return out +} + +func toDriverRegexes(in []regexJSON) []bedrockdriver.GuardrailRegex { + out := make([]bedrockdriver.GuardrailRegex, len(in)) + for i, r := range in { + out[i] = bedrockdriver.GuardrailRegex{Name: r.Name, Pattern: r.Pattern, Action: r.Action, Description: r.Description} + } + + return out +} + +func toDriverGroundingFilters(in []contextualGroundingFilterJSON) []bedrockdriver.GuardrailContextualGroundingFilter { + out := make([]bedrockdriver.GuardrailContextualGroundingFilter, len(in)) + for i, f := range in { + out[i] = bedrockdriver.GuardrailContextualGroundingFilter{Type: f.Type, Threshold: f.Threshold, Action: f.Action} + } + + return out +} + +// --- driver -> response conversion --- + +func fillGuardrailPolicies(out *guardrailJSON, p *bedrockdriver.GuardrailPolicies) { + if tp := p.TopicPolicy; tp != nil { + out.TopicPolicy = &topicPolicyJSON{Topics: fromDriverTopics(tp.Topics)} + } + + if cp := p.ContentPolicy; cp != nil { + out.ContentPolicy = &contentPolicyJSON{Filters: fromDriverContentFilters(cp.Filters)} + } + + if wp := p.WordPolicy; wp != nil { + out.WordPolicy = &wordPolicyJSON{ + Words: fromDriverWords(wp.Words), + ManagedWordLists: fromDriverManagedWords(wp.ManagedWordLists), + } + } + + if sp := p.SensitiveInformationPolicy; sp != nil { + out.SensitiveInformationPolicy = &sensitiveInfoPolicyJSON{ + PiiEntities: fromDriverPiiEntities(sp.PiiEntities), + Regexes: fromDriverRegexes(sp.Regexes), + } + } + + if gp := p.ContextualGroundingPolicy; gp != nil { + out.ContextualGroundingPolicy = &contextualGroundingPolicyJSON{Filters: fromDriverGroundingFilters(gp.Filters)} + } +} + +func fromDriverTopics(in []bedrockdriver.GuardrailTopic) []topicJSON { + out := make([]topicJSON, len(in)) + for i, t := range in { + out[i] = topicJSON{Name: t.Name, Definition: t.Definition, Examples: t.Examples, Type: t.Type} + } + + return out +} + +func fromDriverContentFilters(in []bedrockdriver.GuardrailContentFilter) []contentFilterJSON { + out := make([]contentFilterJSON, len(in)) + for i, f := range in { + out[i] = contentFilterJSON{Type: f.Type, InputStrength: f.InputStrength, OutputStrength: f.OutputStrength} + } + + return out +} + +func fromDriverWords(in []bedrockdriver.GuardrailWord) []wordJSON { + out := make([]wordJSON, len(in)) + for i, w := range in { + out[i] = wordJSON{Text: w.Text} + } + + return out +} + +func fromDriverManagedWords(in []bedrockdriver.GuardrailManagedWordList) []managedWordsJSON { + out := make([]managedWordsJSON, len(in)) + for i, w := range in { + out[i] = managedWordsJSON{Type: w.Type} + } + + return out +} + +func fromDriverPiiEntities(in []bedrockdriver.GuardrailPiiEntity) []piiEntityJSON { + out := make([]piiEntityJSON, len(in)) + for i, e := range in { + out[i] = piiEntityJSON{Type: e.Type, Action: e.Action} + } + + return out +} + +func fromDriverRegexes(in []bedrockdriver.GuardrailRegex) []regexJSON { + out := make([]regexJSON, len(in)) + for i, r := range in { + out[i] = regexJSON{Name: r.Name, Pattern: r.Pattern, Action: r.Action, Description: r.Description} + } + + return out +} + +func fromDriverGroundingFilters(in []bedrockdriver.GuardrailContextualGroundingFilter) []contextualGroundingFilterJSON { + out := make([]contextualGroundingFilterJSON, len(in)) + for i, f := range in { + out[i] = contextualGroundingFilterJSON{Type: f.Type, Threshold: f.Threshold, Action: f.Action} + } + + return out +} diff --git a/server/aws/bedrock/handler.go b/server/aws/bedrock/handler.go index cafa802d..75a32568 100644 --- a/server/aws/bedrock/handler.go +++ b/server/aws/bedrock/handler.go @@ -17,6 +17,11 @@ // DELETE /custom-models/{modelIdentifier} — DeleteCustomModel // POST /model/{modelId}/invoke — InvokeModel // POST /model/{modelId}/converse — Converse +// POST /model/{modelId}/count-tokens — CountTokens +// POST /guardrail/{id}/version/{version}/apply — ApplyGuardrail +// POST /tagResource — TagResource +// POST /untagResource — UntagResource +// POST /listTagsForResource — ListTagsForResource // // The Matches predicate is rooted at these prefixes so it does not shadow the // catch-all S3 handler that may be registered alongside. @@ -43,8 +48,40 @@ const ( pathProvisionedList = "/provisioned-model-throughputs" pathLogging = "/logging/modelinvocations" - actionInvoke = "invoke" - actionConverse = "converse" + pathTagResource = "/tagResource" + pathUntagResource = "/untagResource" + pathListTags = "/listTagsForResource" + + // prefixApplyGuardrail is the singular /guardrail/ runtime prefix, distinct + // from the plural /guardrails control-plane collection. + prefixApplyGuardrail = "/guardrail/" + + prefixAsyncInvoke = "/async-invoke" + prefixImportJobs = "/model-import-jobs" + prefixCopyJobs = "/model-copy-jobs" + prefixEvalJobs = "/evaluation-jobs" + // prefixEvalJobStop is the singular /evaluation-job/ prefix used only by + // StopEvaluationJob (POST /evaluation-job/{id}/stop), distinct from the + // plural /evaluation-jobs collection. + prefixEvalJobStop = "/evaluation-job/" + + prefixInferenceProfiles = "/inference-profiles" + prefixPromptRouters = "/prompt-routers" + prefixARPolicies = "/automated-reasoning-policies" + + prefixMarketplaceEndpoints = "/marketplace-model/endpoints" + prefixListFMAgreementOffers = "/list-foundation-model-agreement-offers" + prefixFMAvailability = "/foundation-model-availability" + pathCreateFMAgreement = "/create-foundation-model-agreement" + pathDeleteFMAgreement = "/delete-foundation-model-agreement" + + suffixRegistration = "/registration" + + actionInvoke = "invoke" + actionConverse = "converse" + actionCountTokens = "count-tokens" + actionConverseStream = "converse-stream" + actionInvokeStream = "invoke-with-response-stream" ) // Handler serves AWS Bedrock restJson1 requests against a Bedrock driver. @@ -63,15 +100,88 @@ func New(drv bedrockdriver.Bedrock) *Handler { //nolint:gochecknoglobals // immutable routing table shared by Matches and ServeHTTP var collectionPrefixes = []string{prefixFoundation, prefixJobs, prefixCustom, prefixGuardrails, prefixProvisioned} +// underPrefix reports whether p equals pre or is a child path of pre. +func underPrefix(p, pre string) bool { + return p == pre || strings.HasPrefix(p, pre+"/") +} + // claims reports whether path p belongs to this handler. func claims(p string) bool { for _, pre := range collectionPrefixes { + if underPrefix(p, pre) { + return true + } + } + + if strings.HasPrefix(p, prefixRuntime) || strings.HasPrefix(p, prefixApplyGuardrail) { + return true + } + + return claimsExtra(p) +} + +// claimsExtra reports whether p belongs to a singleton or predicate-routed +// surface not covered by the collection prefixes. Split from claims to keep +// each function's cyclomatic complexity small. +func claimsExtra(p string) bool { + return p == pathProvisionedList || p == pathLogging || isTagPath(p) || + isAsyncOrJobPath(p) || isRegistryPath(p) || isMarketplaceOrAgreementPath(p) +} + +// marketplaceAgreementPrefixes are the collection/action prefixes for the +// marketplace-model-endpoint and foundation-model-agreement surfaces. +// +//nolint:gochecknoglobals // immutable routing table shared by claims and ServeHTTP +var marketplaceAgreementPrefixes = []string{prefixMarketplaceEndpoints, prefixListFMAgreementOffers, prefixFMAvailability} + +// isMarketplaceOrAgreementPath reports whether p belongs to the +// marketplace-model-endpoint or foundation-model-agreement surface. +func isMarketplaceOrAgreementPath(p string) bool { + if p == pathCreateFMAgreement || p == pathDeleteFMAgreement { + return true + } + + for _, pre := range marketplaceAgreementPrefixes { + if p == pre || strings.HasPrefix(p, pre+"/") { + return true + } + } + + return false +} + +// registryPrefixes are the collection prefixes for the inference-profile, +// prompt-router, and automated-reasoning-policy control-plane surfaces. +// +//nolint:gochecknoglobals // immutable routing table shared by claims and ServeHTTP +var registryPrefixes = []string{prefixInferenceProfiles, prefixPromptRouters, prefixARPolicies} + +// isRegistryPath reports whether p belongs to a control-plane registry surface. +func isRegistryPath(p string) bool { + for _, pre := range registryPrefixes { + if p == pre || strings.HasPrefix(p, pre+"/") { + return true + } + } + + return false +} + +// asyncJobPrefixes are the collection prefixes for the async-invoke and +// control-plane job surfaces. +// +//nolint:gochecknoglobals // immutable routing table shared by claims and ServeHTTP +var asyncJobPrefixes = []string{prefixAsyncInvoke, prefixImportJobs, prefixCopyJobs, prefixEvalJobs} + +// isAsyncOrJobPath reports whether p belongs to the async-invoke or job surface. +func isAsyncOrJobPath(p string) bool { + for _, pre := range asyncJobPrefixes { if p == pre || strings.HasPrefix(p, pre+"/") { return true } } - return strings.HasPrefix(p, prefixRuntime) || p == pathProvisionedList || p == pathLogging + return strings.HasPrefix(p, prefixEvalJobStop) } // Matches claims the Bedrock control-plane and runtime URL prefixes. @@ -84,19 +194,82 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { p := r.URL.Path switch { - case p == prefixFoundation || strings.HasPrefix(p, prefixFoundation+"/"): + case underPrefix(p, prefixFoundation): h.serveFoundation(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixFoundation), "/")) - case p == prefixJobs || strings.HasPrefix(p, prefixJobs+"/"): + case underPrefix(p, prefixJobs): h.serveJobs(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixJobs), "/")) - case p == prefixCustom || strings.HasPrefix(p, prefixCustom+"/"): + case underPrefix(p, prefixCustom): h.serveCustomModels(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixCustom), "/")) case strings.HasPrefix(p, prefixRuntime): h.serveRuntime(w, r, strings.TrimPrefix(p, prefixRuntime)) + case isAsyncOrJobPath(p): + h.serveAsyncJobs(w, r, p) + case isRegistryPath(p): + h.serveRegistries(w, r, p) + case isMarketplaceOrAgreementPath(p): + h.serveMarketplaceAgreements(w, r, p) default: h.serveManagement(w, r, p) } } +// isTagPath reports whether p is one of the resource-tagging endpoints. +func isTagPath(p string) bool { + return p == pathTagResource || p == pathUntagResource || p == pathListTags +} + +// serveTags routes the resource-tagging surface. Each path is POST-only. +func (h *Handler) serveTags(w http.ResponseWriter, r *http.Request, p string) { + if r.Method != http.MethodPost { + methodNotAllowed(w) + + return + } + + switch p { + case pathTagResource: + h.tagResource(w, r) + case pathUntagResource: + h.untagResource(w, r) + case pathListTags: + h.listTagsForResource(w, r) + default: + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported path: "+p) + } +} + +// serveApplyGuardrail handles POST +// /guardrail/{guardrailIdentifier}/version/{guardrailVersion}/apply. rest is the +// path with the /guardrail/ prefix already trimmed. +func (h *Handler) serveApplyGuardrail(w http.ResponseWriter, r *http.Request, rest string) { + // Shape: {guardrailIdentifier}/version/{guardrailVersion}/apply. The + // identifier may be an ARN containing slashes, so anchor on the fixed + // "/version/" separator and the "/apply" suffix rather than a fixed split. + body, ok := strings.CutSuffix(rest, "/apply") + if !ok { + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported guardrail path") + + return + } + + idx := strings.LastIndex(body, "/version/") + if idx < 0 { + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported guardrail path") + + return + } + + identifier, version := body[:idx], body[idx+len("/version/"):] + + if r.Method != http.MethodPost { + methodNotAllowed(w) + + return + } + + h.applyGuardrail(w, r, identifier, version) +} + // serveManagement routes the guardrail, provisioned-throughput, and invocation // logging surfaces. Split out of ServeHTTP to keep each dispatcher small. func (h *Handler) serveManagement(w http.ResponseWriter, r *http.Request, p string) { @@ -109,6 +282,10 @@ func (h *Handler) serveManagement(w http.ResponseWriter, r *http.Request, p stri h.serveGuardrails(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixGuardrails), "/")) case p == pathLogging: h.serveLogging(w, r) + case isTagPath(p): + h.serveTags(w, r, p) + case strings.HasPrefix(p, prefixApplyGuardrail): + h.serveApplyGuardrail(w, r, strings.TrimPrefix(p, prefixApplyGuardrail)) default: writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported path: "+p) } @@ -130,6 +307,8 @@ func (h *Handler) serveGuardrails(w http.ResponseWriter, r *http.Request, id str } switch r.Method { + case http.MethodPost: + h.createGuardrailVersion(w, r, id) case http.MethodGet: h.getGuardrail(w, r, id) case http.MethodPut: @@ -280,6 +459,12 @@ func (h *Handler) serveRuntime(w http.ResponseWriter, r *http.Request, rest stri h.invokeModel(w, r, modelID) case actionConverse: h.converse(w, r, modelID) + case actionConverseStream: + h.converseStream(w, r, modelID) + case actionInvokeStream: + h.invokeModelStream(w, r, modelID) + case actionCountTokens: + h.countTokens(w, r, modelID) default: writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unknown runtime action: "+action) } diff --git a/server/aws/bedrock/jobs.go b/server/aws/bedrock/jobs.go new file mode 100644 index 00000000..791a004c --- /dev/null +++ b/server/aws/bedrock/jobs.go @@ -0,0 +1,517 @@ +package bedrock + +import ( + "encoding/json" + "net/http" + "strings" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- shared wire types --- + +type s3DataSourceJSON struct { + S3URI string `json:"s3Uri,omitempty"` +} + +type modelDataSourceJSON struct { + S3DataSource *s3DataSourceJSON `json:"s3DataSource,omitempty"` +} + +type evalOutputDataConfigJSON struct { + S3URI string `json:"s3Uri,omitempty"` +} + +type createJobArnResponse struct { + JobARN string `json:"jobArn"` +} + +// --- model import job wire types --- + +type createImportJobRequest struct { + JobName string `json:"jobName"` + ImportedModelName string `json:"importedModelName"` + RoleARN string `json:"roleArn"` + ModelDataSource *modelDataSourceJSON `json:"modelDataSource"` + ClientRequestToken string `json:"clientRequestToken"` + ImportedModelKMSKeyID string `json:"importedModelKmsKeyId"` + JobTags []tagPair `json:"jobTags"` + ImportedModelTags []tagPair `json:"importedModelTags"` +} + +type importJobJSON struct { + CreationTime string `json:"creationTime,omitempty"` + EndTime string `json:"endTime,omitempty"` + ImportedModelARN string `json:"importedModelArn,omitempty"` + ImportedModelName string `json:"importedModelName,omitempty"` + JobARN string `json:"jobArn"` + JobName string `json:"jobName"` + ModelDataSource *modelDataSourceJSON `json:"modelDataSource,omitempty"` + RoleARN string `json:"roleArn,omitempty"` + Status string `json:"status"` + LastModifiedTime string `json:"lastModifiedTime,omitempty"` + FailureMessage string `json:"failureMessage,omitempty"` +} + +type importJobSummaryJSON struct { + CreationTime string `json:"creationTime,omitempty"` + EndTime string `json:"endTime,omitempty"` + ImportedModelARN string `json:"importedModelArn,omitempty"` + ImportedModelName string `json:"importedModelName,omitempty"` + JobARN string `json:"jobArn"` + JobName string `json:"jobName"` + Status string `json:"status"` + LastModifiedTime string `json:"lastModifiedTime,omitempty"` +} + +type listImportJobsResponse struct { + ModelImportJobSummaries []importJobSummaryJSON `json:"modelImportJobSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- model copy job wire types --- + +type createCopyJobRequest struct { + SourceModelARN string `json:"sourceModelArn"` + TargetModelName string `json:"targetModelName"` + ClientRequestToken string `json:"clientRequestToken"` + ModelKMSKeyID string `json:"modelKmsKeyId"` + TargetModelTags []tagPair `json:"targetModelTags"` +} + +type copyJobJSON struct { + CreationTime string `json:"creationTime,omitempty"` + JobARN string `json:"jobArn"` + SourceAccountID string `json:"sourceAccountId,omitempty"` + SourceModelARN string `json:"sourceModelArn"` + SourceModelName string `json:"sourceModelName,omitempty"` + Status string `json:"status"` + TargetModelARN string `json:"targetModelArn,omitempty"` + TargetModelName string `json:"targetModelName,omitempty"` + TargetModelKMSKeyARN string `json:"targetModelKmsKeyArn,omitempty"` + FailureMessage string `json:"failureMessage,omitempty"` +} + +type listCopyJobsResponse struct { + ModelCopyJobSummaries []copyJobJSON `json:"modelCopyJobSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- evaluation job wire types --- + +type createEvalJobRequest struct { + JobName string `json:"jobName"` + RoleARN string `json:"roleArn"` + EvaluationConfig json.RawMessage `json:"evaluationConfig"` + InferenceConfig json.RawMessage `json:"inferenceConfig"` + OutputDataConfig *evalOutputDataConfigJSON `json:"outputDataConfig"` + ApplicationType string `json:"applicationType"` + ClientRequestToken string `json:"clientRequestToken"` + CustomerEncryptionKeyID string `json:"customerEncryptionKeyId"` + JobDescription string `json:"jobDescription"` + JobTags []tagPair `json:"jobTags"` +} + +type evalJobJSON struct { + ApplicationType string `json:"applicationType,omitempty"` + CreationTime string `json:"creationTime,omitempty"` + EvaluationConfig json.RawMessage `json:"evaluationConfig,omitempty"` + InferenceConfig json.RawMessage `json:"inferenceConfig,omitempty"` + JobARN string `json:"jobArn"` + JobName string `json:"jobName"` + JobType string `json:"jobType,omitempty"` + OutputDataConfig *evalOutputDataConfigJSON `json:"outputDataConfig,omitempty"` + RoleARN string `json:"roleArn,omitempty"` + Status string `json:"status"` + LastModifiedTime string `json:"lastModifiedTime,omitempty"` + JobDescription string `json:"jobDescription,omitempty"` + CustomerEncryptionKeyID string `json:"customerEncryptionKeyId,omitempty"` + FailureMessages []string `json:"failureMessages,omitempty"` +} + +type evalJobSummaryJSON struct { + ApplicationType string `json:"applicationType,omitempty"` + CreationTime string `json:"creationTime,omitempty"` + JobARN string `json:"jobArn"` + JobName string `json:"jobName"` + JobType string `json:"jobType"` + Status string `json:"status"` + EvaluationTaskTypes []string `json:"evaluationTaskTypes,omitempty"` +} + +type listEvalJobsResponse struct { + JobSummaries []evalJobSummaryJSON `json:"jobSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- import job dispatch + operations --- + +// serveImportJobs handles /model-import-jobs[/{jobIdentifier}]. +func (h *Handler) serveImportJobs(w http.ResponseWriter, r *http.Request, id string) { + if id == "" { + switch r.Method { + case http.MethodPost: + h.createImportJob(w, r) + case http.MethodGet: + h.listImportJobs(w, r) + default: + methodNotAllowed(w) + } + + return + } + + if r.Method != http.MethodGet { + methodNotAllowed(w) + + return + } + + h.getImportJob(w, r, id) +} + +func (h *Handler) createImportJob(w http.ResponseWriter, r *http.Request) { + var in createImportJobRequest + if !decodeJSON(w, r, &in) { + return + } + + job, err := h.bedrock.CreateModelImportJob(r.Context(), bedrockdriver.ModelImportJobConfig{ + JobName: in.JobName, + ImportedModelName: in.ImportedModelName, + RoleARN: in.RoleARN, + ModelDataSourceS3URI: modelDataSourceURI(in.ModelDataSource), + ClientRequestToken: in.ClientRequestToken, + ImportedModelKMSKeyID: in.ImportedModelKMSKeyID, + JobTags: tagsToMap(in.JobTags), + ImportedModelTags: tagsToMap(in.ImportedModelTags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createJobArnResponse{JobARN: job.JobARN}) +} + +func (h *Handler) getImportJob(w http.ResponseWriter, r *http.Request, id string) { + job, err := h.bedrock.GetModelImportJob(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toImportJobJSON(job)) +} + +func (h *Handler) listImportJobs(w http.ResponseWriter, r *http.Request) { + jobs, err := h.bedrock.ListModelImportJobs(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]importJobSummaryJSON, 0, len(jobs)) + for i := range jobs { + out = append(out, toImportJobSummaryJSON(&jobs[i])) + } + + writeJSON(w, listImportJobsResponse{ModelImportJobSummaries: out}) +} + +// --- copy job dispatch + operations --- + +// serveCopyJobs handles /model-copy-jobs[/{jobArn}]. The job ARN contains +// slashes, so it is the entire remainder of the path. +func (h *Handler) serveCopyJobs(w http.ResponseWriter, r *http.Request, arn string) { + if arn == "" { + switch r.Method { + case http.MethodPost: + h.createCopyJob(w, r) + case http.MethodGet: + h.listCopyJobs(w, r) + default: + methodNotAllowed(w) + } + + return + } + + if r.Method != http.MethodGet { + methodNotAllowed(w) + + return + } + + h.getCopyJob(w, r, arn) +} + +func (h *Handler) createCopyJob(w http.ResponseWriter, r *http.Request) { + var in createCopyJobRequest + if !decodeJSON(w, r, &in) { + return + } + + job, err := h.bedrock.CreateModelCopyJob(r.Context(), bedrockdriver.ModelCopyJobConfig{ + SourceModelARN: in.SourceModelARN, + TargetModelName: in.TargetModelName, + ClientRequestToken: in.ClientRequestToken, + ModelKMSKeyID: in.ModelKMSKeyID, + TargetModelTags: tagsToMap(in.TargetModelTags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createJobArnResponse{JobARN: job.JobARN}) +} + +func (h *Handler) getCopyJob(w http.ResponseWriter, r *http.Request, arn string) { + job, err := h.bedrock.GetModelCopyJob(r.Context(), arn) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toCopyJobJSON(job)) +} + +func (h *Handler) listCopyJobs(w http.ResponseWriter, r *http.Request) { + jobs, err := h.bedrock.ListModelCopyJobs(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]copyJobJSON, 0, len(jobs)) + for i := range jobs { + out = append(out, toCopyJobJSON(&jobs[i])) + } + + writeJSON(w, listCopyJobsResponse{ModelCopyJobSummaries: out}) +} + +// --- evaluation job dispatch + operations --- + +// serveEvalJobs handles /evaluation-jobs[/{jobIdentifier}]. +func (h *Handler) serveEvalJobs(w http.ResponseWriter, r *http.Request, id string) { + if id == "" { + switch r.Method { + case http.MethodPost: + h.createEvalJob(w, r) + case http.MethodGet: + h.listEvalJobs(w, r) + default: + methodNotAllowed(w) + } + + return + } + + if r.Method != http.MethodGet { + methodNotAllowed(w) + + return + } + + h.getEvalJob(w, r, id) +} + +// serveEvalJobStop handles POST /evaluation-job/{jobIdentifier}/stop. rest is +// the path with the /evaluation-job/ prefix already trimmed; the identifier may +// contain slashes (ARN), so /stop is split off the tail. +func (h *Handler) serveEvalJobStop(w http.ResponseWriter, r *http.Request, rest string) { + const suffix = "/stop" + + if !strings.HasSuffix(rest, suffix) { + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported evaluation-job path") + + return + } + + id := strings.TrimSuffix(rest, suffix) + + if r.Method != http.MethodPost { + methodNotAllowed(w) + + return + } + + if err := h.bedrock.StopEvaluationJob(r.Context(), id); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +func (h *Handler) createEvalJob(w http.ResponseWriter, r *http.Request) { + var in createEvalJobRequest + if !decodeJSON(w, r, &in) { + return + } + + job, err := h.bedrock.CreateEvaluationJob(r.Context(), bedrockdriver.EvaluationJobConfig{ + JobName: in.JobName, + RoleARN: in.RoleARN, + EvaluationConfig: []byte(in.EvaluationConfig), + InferenceConfig: []byte(in.InferenceConfig), + OutputDataS3URI: evalOutputURI(in.OutputDataConfig), + ApplicationType: in.ApplicationType, + ClientRequestToken: in.ClientRequestToken, + CustomerEncryptionKeyID: in.CustomerEncryptionKeyID, + JobDescription: in.JobDescription, + JobTags: tagsToMap(in.JobTags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createJobArnResponse{JobARN: job.JobARN}) +} + +func (h *Handler) getEvalJob(w http.ResponseWriter, r *http.Request, id string) { + job, err := h.bedrock.GetEvaluationJob(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toEvalJobJSON(job)) +} + +func (h *Handler) listEvalJobs(w http.ResponseWriter, r *http.Request) { + jobs, err := h.bedrock.ListEvaluationJobs(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]evalJobSummaryJSON, 0, len(jobs)) + for i := range jobs { + out = append(out, toEvalJobSummaryJSON(&jobs[i])) + } + + writeJSON(w, listEvalJobsResponse{JobSummaries: out}) +} + +// --- converters --- + +func modelDataSourceURI(in *modelDataSourceJSON) string { + if in == nil || in.S3DataSource == nil { + return "" + } + + return in.S3DataSource.S3URI +} + +func evalOutputURI(in *evalOutputDataConfigJSON) string { + if in == nil { + return "" + } + + return in.S3URI +} + +func dataSourceOf(uri string) *modelDataSourceJSON { + if uri == "" { + return nil + } + + return &modelDataSourceJSON{S3DataSource: &s3DataSourceJSON{S3URI: uri}} +} + +func evalOutputOf(uri string) *evalOutputDataConfigJSON { + if uri == "" { + return nil + } + + return &evalOutputDataConfigJSON{S3URI: uri} +} + +func toImportJobJSON(job *bedrockdriver.ModelImportJob) importJobJSON { + return importJobJSON{ + CreationTime: job.CreationTime, + EndTime: job.EndTime, + ImportedModelARN: job.ImportedModelARN, + ImportedModelName: job.ImportedModelName, + JobARN: job.JobARN, + JobName: job.JobName, + ModelDataSource: dataSourceOf(job.ModelDataSourceS3URI), + RoleARN: job.RoleARN, + Status: job.Status, + LastModifiedTime: job.LastModifiedTime, + FailureMessage: job.FailureMessage, + } +} + +func toImportJobSummaryJSON(job *bedrockdriver.ModelImportJob) importJobSummaryJSON { + return importJobSummaryJSON{ + CreationTime: job.CreationTime, + EndTime: job.EndTime, + ImportedModelARN: job.ImportedModelARN, + ImportedModelName: job.ImportedModelName, + JobARN: job.JobARN, + JobName: job.JobName, + Status: job.Status, + LastModifiedTime: job.LastModifiedTime, + } +} + +func toCopyJobJSON(job *bedrockdriver.ModelCopyJob) copyJobJSON { + return copyJobJSON{ + CreationTime: job.CreationTime, + JobARN: job.JobARN, + SourceAccountID: job.SourceAccountID, + SourceModelARN: job.SourceModelARN, + SourceModelName: job.SourceModelName, + Status: job.Status, + TargetModelARN: job.TargetModelARN, + TargetModelName: job.TargetModelName, + TargetModelKMSKeyARN: job.TargetModelKMSKeyARN, + FailureMessage: job.FailureMessage, + } +} + +func toEvalJobJSON(job *bedrockdriver.EvaluationJob) evalJobJSON { + return evalJobJSON{ + ApplicationType: job.ApplicationType, + CreationTime: job.CreationTime, + EvaluationConfig: json.RawMessage(job.EvaluationConfig), + InferenceConfig: json.RawMessage(job.InferenceConfig), + JobARN: job.JobARN, + JobName: job.JobName, + JobType: job.JobType, + OutputDataConfig: evalOutputOf(job.OutputDataS3URI), + RoleARN: job.RoleARN, + Status: job.Status, + LastModifiedTime: job.LastModifiedTime, + JobDescription: job.JobDescription, + CustomerEncryptionKeyID: job.CustomerEncryptionKeyID, + FailureMessages: job.FailureMessages, + } +} + +func toEvalJobSummaryJSON(job *bedrockdriver.EvaluationJob) evalJobSummaryJSON { + return evalJobSummaryJSON{ + ApplicationType: job.ApplicationType, + CreationTime: job.CreationTime, + JobARN: job.JobARN, + JobName: job.JobName, + JobType: job.JobType, + Status: job.Status, + EvaluationTaskTypes: []string{"Generation"}, + } +} diff --git a/server/aws/bedrock/management.go b/server/aws/bedrock/management.go index d055b35b..1e529943 100644 --- a/server/aws/bedrock/management.go +++ b/server/aws/bedrock/management.go @@ -2,6 +2,7 @@ package bedrock import ( "net/http" + "strconv" bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" ) @@ -14,13 +15,28 @@ type tagPair struct { } type createGuardrailRequest struct { - Name string `json:"name"` - Description string `json:"description"` - BlockedInputMessaging string `json:"blockedInputMessaging"` - BlockedOutputsMessaging string `json:"blockedOutputsMessaging"` - KMSKeyID string `json:"kmsKeyId"` - ClientRequestToken string `json:"clientRequestToken"` - Tags []tagPair `json:"tags"` + Name string `json:"name"` + Description string `json:"description"` + BlockedInputMessaging string `json:"blockedInputMessaging"` + BlockedOutputsMessaging string `json:"blockedOutputsMessaging"` + KMSKeyID string `json:"kmsKeyId"` + ClientRequestToken string `json:"clientRequestToken"` + Tags []tagPair `json:"tags"` + TopicPolicyConfig *topicPolicyConfigJSON `json:"topicPolicyConfig,omitempty"` + ContentPolicyConfig *contentPolicyConfigJSON `json:"contentPolicyConfig,omitempty"` + WordPolicyConfig *wordPolicyConfigJSON `json:"wordPolicyConfig,omitempty"` + SensitiveInformationPolicyConfig *sensitiveInfoPolicyConfigJSON `json:"sensitiveInformationPolicyConfig,omitempty"` + ContextualGroundingPolicyConfig *contextualGroundingPolicyConfigJSON `json:"contextualGroundingPolicyConfig,omitempty"` +} + +type createGuardrailVersionRequest struct { + ClientRequestToken string `json:"clientRequestToken"` + Description string `json:"description"` +} + +type createGuardrailVersionResponse struct { + GuardrailID string `json:"guardrailId"` + Version string `json:"version"` } type createGuardrailResponse struct { @@ -31,17 +47,22 @@ type createGuardrailResponse struct { } type guardrailJSON struct { - Name string `json:"name"` - Description string `json:"description,omitempty"` - GuardrailID string `json:"guardrailId"` - GuardrailARN string `json:"guardrailArn"` - Version string `json:"version"` - Status string `json:"status"` - BlockedInputMessaging string `json:"blockedInputMessaging,omitempty"` - BlockedOutputsMessaging string `json:"blockedOutputsMessaging,omitempty"` - KMSKeyARN string `json:"kmsKeyArn,omitempty"` - CreatedAt string `json:"createdAt,omitempty"` - UpdatedAt string `json:"updatedAt,omitempty"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + GuardrailID string `json:"guardrailId"` + GuardrailARN string `json:"guardrailArn"` + Version string `json:"version"` + Status string `json:"status"` + BlockedInputMessaging string `json:"blockedInputMessaging,omitempty"` + BlockedOutputsMessaging string `json:"blockedOutputsMessaging,omitempty"` + KMSKeyARN string `json:"kmsKeyArn,omitempty"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` + TopicPolicy *topicPolicyJSON `json:"topicPolicy,omitempty"` + ContentPolicy *contentPolicyJSON `json:"contentPolicy,omitempty"` + WordPolicy *wordPolicyJSON `json:"wordPolicy,omitempty"` + SensitiveInformationPolicy *sensitiveInfoPolicyJSON `json:"sensitiveInformationPolicy,omitempty"` + ContextualGroundingPolicy *contextualGroundingPolicyJSON `json:"contextualGroundingPolicy,omitempty"` } type guardrailSummaryJSON struct { @@ -143,6 +164,7 @@ func (h *Handler) createGuardrail(w http.ResponseWriter, r *http.Request) { KMSKeyID: in.KMSKeyID, ClientRequestToken: in.ClientRequestToken, Tags: tagsToMap(in.Tags), + GuardrailPolicies: toDriverGuardrailPolicies(&in), }) if err != nil { writeErr(w, err) @@ -155,6 +177,22 @@ func (h *Handler) createGuardrail(w http.ResponseWriter, r *http.Request) { }) } +func (h *Handler) createGuardrailVersion(w http.ResponseWriter, r *http.Request, id string) { + var in createGuardrailVersionRequest + if !decodeJSON(w, r, &in) { + return + } + + gid, version, err := h.bedrock.CreateGuardrailVersion(r.Context(), id, in.Description) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createGuardrailVersionResponse{GuardrailID: gid, Version: version}) +} + func (h *Handler) getGuardrail(w http.ResponseWriter, r *http.Request, id string) { g, err := h.bedrock.GetGuardrail(r.Context(), id, r.URL.Query().Get("guardrailVersion")) if err != nil { @@ -167,19 +205,43 @@ func (h *Handler) getGuardrail(w http.ResponseWriter, r *http.Request, id string } func (h *Handler) listGuardrails(w http.ResponseWriter, r *http.Request) { - gs, err := h.bedrock.ListGuardrails(r.Context()) + q := r.URL.Query() + + gs, err := h.bedrock.ListGuardrails(r.Context(), q.Get("guardrailIdentifier")) if err != nil { writeErr(w, err) return } - out := make([]guardrailSummaryJSON, 0, len(gs)) - for i := range gs { - out = append(out, toGuardrailSummaryJSON(&gs[i])) + page, next := paginateGuardrails(gs, q.Get("maxResults"), q.Get("nextToken")) + + out := make([]guardrailSummaryJSON, 0, len(page)) + for i := range page { + out = append(out, toGuardrailSummaryJSON(&page[i])) + } + + writeJSON(w, listGuardrailsResponse{Guardrails: out, NextToken: next}) +} + +// paginateGuardrails applies optional maxResults/nextToken (a decimal start +// offset) paging, returning the page and the token for the next page (empty +// when the page is the last). +func paginateGuardrails( + gs []bedrockdriver.Guardrail, maxResults, nextToken string, +) (page []bedrockdriver.Guardrail, next string) { + start := 0 + if n, err := strconv.Atoi(nextToken); err == nil && n > 0 { + start = min(n, len(gs)) + } + + gs = gs[start:] + + if n, err := strconv.Atoi(maxResults); err == nil && n > 0 && n < len(gs) { + return gs[:n], strconv.Itoa(start + n) } - writeJSON(w, listGuardrailsResponse{Guardrails: out}) + return gs, "" } func (h *Handler) updateGuardrail(w http.ResponseWriter, r *http.Request, id string) { @@ -194,6 +256,7 @@ func (h *Handler) updateGuardrail(w http.ResponseWriter, r *http.Request, id str BlockedInputMessaging: in.BlockedInputMessaging, BlockedOutputsMessaging: in.BlockedOutputsMessaging, KMSKeyID: in.KMSKeyID, + GuardrailPolicies: toDriverGuardrailPolicies(&in), }) if err != nil { writeErr(w, err) @@ -207,7 +270,7 @@ func (h *Handler) updateGuardrail(w http.ResponseWriter, r *http.Request, id str } func (h *Handler) deleteGuardrail(w http.ResponseWriter, r *http.Request, id string) { - if err := h.bedrock.DeleteGuardrail(r.Context(), id); err != nil { + if err := h.bedrock.DeleteGuardrail(r.Context(), id, r.URL.Query().Get("guardrailVersion")); err != nil { writeErr(w, err) return @@ -326,7 +389,7 @@ func (h *Handler) deleteLogging(w http.ResponseWriter, r *http.Request) { // --- converters --- func toGuardrailJSON(g *bedrockdriver.Guardrail) guardrailJSON { - return guardrailJSON{ + out := guardrailJSON{ Name: g.Name, Description: g.Description, GuardrailID: g.ID, @@ -339,6 +402,10 @@ func toGuardrailJSON(g *bedrockdriver.Guardrail) guardrailJSON { CreatedAt: g.CreatedAt, UpdatedAt: g.UpdatedAt, } + + fillGuardrailPolicies(&out, &g.GuardrailPolicies) + + return out } func toGuardrailSummaryJSON(g *bedrockdriver.Guardrail) guardrailSummaryJSON { diff --git a/server/aws/bedrock/marketplace_agreements.go b/server/aws/bedrock/marketplace_agreements.go new file mode 100644 index 00000000..e4f335a7 --- /dev/null +++ b/server/aws/bedrock/marketplace_agreements.go @@ -0,0 +1,400 @@ +package bedrock + +import ( + "encoding/json" + "net/http" + "strings" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- marketplace model endpoint wire types --- + +type createMarketplaceEndpointRequest struct { + EndpointConfig json.RawMessage `json:"endpointConfig"` + EndpointName string `json:"endpointName"` + ModelSourceIdentifier string `json:"modelSourceIdentifier"` + AcceptEula bool `json:"acceptEula"` + ClientRequestToken string `json:"clientRequestToken"` + Tags []tagPair `json:"tags"` +} + +type updateMarketplaceEndpointRequest struct { + EndpointConfig json.RawMessage `json:"endpointConfig"` + ClientRequestToken string `json:"clientRequestToken"` +} + +type registerMarketplaceEndpointRequest struct { + ModelSourceIdentifier string `json:"modelSourceIdentifier"` +} + +type marketplaceEndpointJSON struct { + EndpointARN string `json:"endpointArn"` + ModelSourceIdentifier string `json:"modelSourceIdentifier,omitempty"` + EndpointConfig json.RawMessage `json:"endpointConfig,omitempty"` + EndpointStatus string `json:"endpointStatus,omitempty"` + Status string `json:"status,omitempty"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` + EndpointStatusMessage string `json:"endpointStatusMessage,omitempty"` + StatusMessage string `json:"statusMessage,omitempty"` +} + +type marketplaceEndpointSummaryJSON struct { + EndpointARN string `json:"endpointArn"` + ModelSourceIdentifier string `json:"modelSourceIdentifier,omitempty"` + Status string `json:"status,omitempty"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` +} + +type marketplaceEndpointResponse struct { + MarketplaceModelEndpoint marketplaceEndpointJSON `json:"marketplaceModelEndpoint"` +} + +type listMarketplaceEndpointsResponse struct { + MarketplaceModelEndpoints []marketplaceEndpointSummaryJSON `json:"marketplaceModelEndpoints"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- foundation model agreement wire types --- + +type createFMAgreementRequest struct { + ModelID string `json:"modelId"` + OfferToken string `json:"offerToken"` +} + +type createFMAgreementResponse struct { + ModelID string `json:"modelId,omitempty"` +} + +type deleteFMAgreementRequest struct { + ModelID string `json:"modelId"` +} + +// fmOfferTermDetailsJSON is emitted as an empty object; the emulator returns no +// term details for synthetic offers. +type fmOfferTermDetailsJSON struct{} + +type fmOfferJSON struct { + OfferToken string `json:"offerToken,omitempty"` + OfferID string `json:"offerId,omitempty"` + TermDetails fmOfferTermDetailsJSON `json:"termDetails"` +} + +type listFMAgreementOffersResponse struct { + ModelID string `json:"modelId,omitempty"` + Offers []fmOfferJSON `json:"offers"` +} + +type agreementAvailabilityJSON struct { + Status string `json:"status,omitempty"` +} + +type getFMAvailabilityResponse struct { + ModelID string `json:"modelId,omitempty"` + AgreementAvailability *agreementAvailabilityJSON `json:"agreementAvailability,omitempty"` + AuthorizationStatus string `json:"authorizationStatus,omitempty"` + EntitlementAvailability string `json:"entitlementAvailability,omitempty"` + RegionAvailability string `json:"regionAvailability,omitempty"` +} + +// --- dispatcher --- + +// serveMarketplaceAgreements routes the marketplace-model-endpoint and +// foundation-model-agreement control-plane surfaces. Split out of ServeHTTP to +// keep each dispatcher small. +func (h *Handler) serveMarketplaceAgreements(w http.ResponseWriter, r *http.Request, p string) { + switch { + case p == pathCreateFMAgreement: + h.createFoundationModelAgreement(w, r) + case p == pathDeleteFMAgreement: + h.deleteFoundationModelAgreement(w, r) + case p == prefixListFMAgreementOffers || strings.HasPrefix(p, prefixListFMAgreementOffers+"/"): + h.listFoundationModelAgreementOffers(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixListFMAgreementOffers), "/")) + case p == prefixFMAvailability || strings.HasPrefix(p, prefixFMAvailability+"/"): + h.getFoundationModelAvailability(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixFMAvailability), "/")) + default: + h.serveMarketplaceEndpoints(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixMarketplaceEndpoints), "/")) + } +} + +// --- marketplace endpoint dispatch + operations --- + +// serveMarketplaceEndpoints handles /marketplace-model/endpoints and its +// resource + registration sub-paths. The endpoint ARN contains slashes, so it +// is the entire remainder unless it carries the /registration suffix. +func (h *Handler) serveMarketplaceEndpoints(w http.ResponseWriter, r *http.Request, rest string) { + switch { + case rest == "": + h.marketplaceEndpointCollection(w, r) + case strings.HasSuffix(rest, suffixRegistration): + h.marketplaceEndpointRegistration(w, r, strings.TrimSuffix(rest, suffixRegistration)) + default: + h.marketplaceEndpointResource(w, r, rest) + } +} + +// marketplaceEndpointCollection handles POST (create) and GET (list) on the +// /marketplace-model/endpoints collection. +func (h *Handler) marketplaceEndpointCollection(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + h.createMarketplaceEndpoint(w, r) + case http.MethodGet: + h.listMarketplaceEndpoints(w, r) + default: + methodNotAllowed(w) + } +} + +// marketplaceEndpointResource handles GET/PATCH/DELETE on a single endpoint ARN. +func (h *Handler) marketplaceEndpointResource(w http.ResponseWriter, r *http.Request, arn string) { + switch r.Method { + case http.MethodGet: + h.getMarketplaceEndpoint(w, r, arn) + case http.MethodPatch: + h.updateMarketplaceEndpoint(w, r, arn) + case http.MethodDelete: + h.deleteMarketplaceEndpoint(w, r, arn) + default: + methodNotAllowed(w) + } +} + +// marketplaceEndpointRegistration handles POST (register) and DELETE +// (deregister) on {endpointIdentifier}/registration. +func (h *Handler) marketplaceEndpointRegistration(w http.ResponseWriter, r *http.Request, id string) { + switch r.Method { + case http.MethodPost: + h.registerMarketplaceEndpoint(w, r, id) + case http.MethodDelete: + h.deregisterMarketplaceEndpoint(w, r, id) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) createMarketplaceEndpoint(w http.ResponseWriter, r *http.Request) { + var in createMarketplaceEndpointRequest + if !decodeJSON(w, r, &in) { + return + } + + endpoint, err := h.bedrock.CreateMarketplaceModelEndpoint(r.Context(), bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: in.EndpointName, + ModelSourceIdentifier: in.ModelSourceIdentifier, + EndpointConfig: []byte(in.EndpointConfig), + AcceptEula: in.AcceptEula, + ClientRequestToken: in.ClientRequestToken, + Tags: tagsToMap(in.Tags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, marketplaceEndpointResponse{MarketplaceModelEndpoint: toMarketplaceEndpointJSON(endpoint)}) +} + +func (h *Handler) getMarketplaceEndpoint(w http.ResponseWriter, r *http.Request, arn string) { + endpoint, err := h.bedrock.GetMarketplaceModelEndpoint(r.Context(), arn) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, marketplaceEndpointResponse{MarketplaceModelEndpoint: toMarketplaceEndpointJSON(endpoint)}) +} + +func (h *Handler) listMarketplaceEndpoints(w http.ResponseWriter, r *http.Request) { + endpoints, err := h.bedrock.ListMarketplaceModelEndpoints(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + filter := r.URL.Query().Get("modelSourceIdentifier") + out := make([]marketplaceEndpointSummaryJSON, 0, len(endpoints)) + + for i := range endpoints { + if filter != "" && endpoints[i].ModelSourceIdentifier != filter { + continue + } + + out = append(out, toMarketplaceEndpointSummaryJSON(&endpoints[i])) + } + + writeJSON(w, listMarketplaceEndpointsResponse{MarketplaceModelEndpoints: out}) +} + +func (h *Handler) updateMarketplaceEndpoint(w http.ResponseWriter, r *http.Request, arn string) { + var in updateMarketplaceEndpointRequest + if !decodeJSON(w, r, &in) { + return + } + + endpoint, err := h.bedrock.UpdateMarketplaceModelEndpoint(r.Context(), arn, []byte(in.EndpointConfig)) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, marketplaceEndpointResponse{MarketplaceModelEndpoint: toMarketplaceEndpointJSON(endpoint)}) +} + +func (h *Handler) deleteMarketplaceEndpoint(w http.ResponseWriter, r *http.Request, arn string) { + if err := h.bedrock.DeleteMarketplaceModelEndpoint(r.Context(), arn); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +func (h *Handler) registerMarketplaceEndpoint(w http.ResponseWriter, r *http.Request, id string) { + var in registerMarketplaceEndpointRequest + if !decodeJSON(w, r, &in) { + return + } + + endpoint, err := h.bedrock.RegisterMarketplaceModelEndpoint(r.Context(), id, in.ModelSourceIdentifier) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, marketplaceEndpointResponse{MarketplaceModelEndpoint: toMarketplaceEndpointJSON(endpoint)}) +} + +func (h *Handler) deregisterMarketplaceEndpoint(w http.ResponseWriter, r *http.Request, id string) { + if err := h.bedrock.DeregisterMarketplaceModelEndpoint(r.Context(), id); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +// --- foundation model agreement operations --- + +func (h *Handler) createFoundationModelAgreement(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + methodNotAllowed(w) + + return + } + + var in createFMAgreementRequest + if !decodeJSON(w, r, &in) { + return + } + + modelID, err := h.bedrock.CreateFoundationModelAgreement(r.Context(), in.ModelID, in.OfferToken) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createFMAgreementResponse{ModelID: modelID}) +} + +func (h *Handler) deleteFoundationModelAgreement(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + methodNotAllowed(w) + + return + } + + var in deleteFMAgreementRequest + if !decodeJSON(w, r, &in) { + return + } + + if err := h.bedrock.DeleteFoundationModelAgreement(r.Context(), in.ModelID); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +func (h *Handler) listFoundationModelAgreementOffers(w http.ResponseWriter, r *http.Request, modelID string) { + if r.Method != http.MethodGet { + methodNotAllowed(w) + + return + } + + offers, err := h.bedrock.ListFoundationModelAgreementOffers(r.Context(), modelID, r.URL.Query().Get("offerType")) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]fmOfferJSON, 0, len(offers)) + for i := range offers { + out = append(out, fmOfferJSON{OfferToken: offers[i].OfferToken, OfferID: offers[i].OfferID}) + } + + writeJSON(w, listFMAgreementOffersResponse{ModelID: modelID, Offers: out}) +} + +func (h *Handler) getFoundationModelAvailability(w http.ResponseWriter, r *http.Request, modelID string) { + if r.Method != http.MethodGet { + methodNotAllowed(w) + + return + } + + avail, err := h.bedrock.GetFoundationModelAvailability(r.Context(), modelID) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, getFMAvailabilityResponse{ + ModelID: modelID, + AgreementAvailability: &agreementAvailabilityJSON{Status: avail.AgreementStatus}, + AuthorizationStatus: avail.AuthorizationStatus, + EntitlementAvailability: avail.EntitlementAvailability, + RegionAvailability: avail.RegionAvailability, + }) +} + +// --- converters --- + +func toMarketplaceEndpointJSON(e *bedrockdriver.MarketplaceEndpoint) marketplaceEndpointJSON { + return marketplaceEndpointJSON{ + EndpointARN: e.EndpointARN, + ModelSourceIdentifier: e.ModelSourceIdentifier, + EndpointConfig: json.RawMessage(e.EndpointConfig), + EndpointStatus: e.EndpointStatus, + Status: e.Status, + CreatedAt: e.CreatedAt, + UpdatedAt: e.UpdatedAt, + EndpointStatusMessage: e.EndpointStatusMessage, + StatusMessage: e.StatusMessage, + } +} + +func toMarketplaceEndpointSummaryJSON(e *bedrockdriver.MarketplaceEndpoint) marketplaceEndpointSummaryJSON { + return marketplaceEndpointSummaryJSON{ + EndpointARN: e.EndpointARN, + ModelSourceIdentifier: e.ModelSourceIdentifier, + Status: e.Status, + CreatedAt: e.CreatedAt, + UpdatedAt: e.UpdatedAt, + } +} diff --git a/server/aws/bedrock/registries.go b/server/aws/bedrock/registries.go new file mode 100644 index 00000000..a6639b86 --- /dev/null +++ b/server/aws/bedrock/registries.go @@ -0,0 +1,583 @@ +package bedrock + +import ( + "encoding/json" + "net/http" + "strings" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- inference profile wire types --- + +type inferenceProfileModelJSON struct { + ModelARN string `json:"modelArn,omitempty"` +} + +type inferenceProfileModelSourceJSON struct { + CopyFrom string `json:"copyFrom,omitempty"` +} + +type createInferenceProfileRequest struct { + InferenceProfileName string `json:"inferenceProfileName"` + ModelSource *inferenceProfileModelSourceJSON `json:"modelSource"` + ClientRequestToken string `json:"clientRequestToken"` + Description string `json:"description"` + Tags []tagPair `json:"tags"` +} + +type createInferenceProfileResponse struct { + InferenceProfileARN string `json:"inferenceProfileArn"` + Status string `json:"status"` +} + +type inferenceProfileJSON struct { + InferenceProfileARN string `json:"inferenceProfileArn"` + InferenceProfileID string `json:"inferenceProfileId,omitempty"` + InferenceProfileName string `json:"inferenceProfileName,omitempty"` + Models []inferenceProfileModelJSON `json:"models,omitempty"` + Status string `json:"status"` + Type string `json:"type"` + Description string `json:"description,omitempty"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` +} + +type listInferenceProfilesResponse struct { + InferenceProfileSummaries []inferenceProfileJSON `json:"inferenceProfileSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- prompt router wire types --- + +type promptRouterTargetModelJSON struct { + ModelARN string `json:"modelArn,omitempty"` +} + +type routingCriteriaJSON struct { + ResponseQualityDifference *float64 `json:"responseQualityDifference,omitempty"` +} + +type createPromptRouterRequest struct { + PromptRouterName string `json:"promptRouterName"` + Models []promptRouterTargetModelJSON `json:"models"` + RoutingCriteria *routingCriteriaJSON `json:"routingCriteria"` + FallbackModel *promptRouterTargetModelJSON `json:"fallbackModel"` + ClientRequestToken string `json:"clientRequestToken"` + Description string `json:"description"` + Tags []tagPair `json:"tags"` +} + +type createPromptRouterResponse struct { + PromptRouterARN string `json:"promptRouterArn"` +} + +type promptRouterJSON struct { + PromptRouterARN string `json:"promptRouterArn"` + PromptRouterName string `json:"promptRouterName,omitempty"` + Models []promptRouterTargetModelJSON `json:"models,omitempty"` + RoutingCriteria *routingCriteriaJSON `json:"routingCriteria,omitempty"` + FallbackModel *promptRouterTargetModelJSON `json:"fallbackModel,omitempty"` + Status string `json:"status"` + Type string `json:"type"` + Description string `json:"description,omitempty"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` +} + +type listPromptRoutersResponse struct { + PromptRouterSummaries []promptRouterJSON `json:"promptRouterSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- automated reasoning policy wire types --- + +type createARPolicyRequest struct { + Name string `json:"name"` + ClientRequestToken string `json:"clientRequestToken"` + Description string `json:"description"` + KMSKeyID string `json:"kmsKeyId"` + PolicyDefinition json.RawMessage `json:"policyDefinition"` + Tags []tagPair `json:"tags"` +} + +type updateARPolicyRequest struct { + PolicyDefinition json.RawMessage `json:"policyDefinition"` + Description string `json:"description"` + Name string `json:"name"` +} + +type createARPolicyResponse struct { + PolicyARN string `json:"policyArn"` + Name string `json:"name"` + Version string `json:"version"` + DefinitionHash string `json:"definitionHash"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` + Description string `json:"description,omitempty"` +} + +type arPolicyJSON struct { + PolicyARN string `json:"policyArn"` + PolicyID string `json:"policyId,omitempty"` + Name string `json:"name"` + Version string `json:"version"` + DefinitionHash string `json:"definitionHash"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` + Description string `json:"description,omitempty"` + KMSKeyARN string `json:"kmsKeyArn,omitempty"` +} + +type arPolicySummaryJSON struct { + PolicyARN string `json:"policyArn"` + PolicyID string `json:"policyId,omitempty"` + Name string `json:"name"` + Version string `json:"version"` + CreatedAt string `json:"createdAt,omitempty"` + UpdatedAt string `json:"updatedAt,omitempty"` + Description string `json:"description,omitempty"` +} + +type listARPoliciesResponse struct { + AutomatedReasoningPolicySummaries []arPolicySummaryJSON `json:"automatedReasoningPolicySummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +type updateARPolicyResponse struct { + PolicyARN string `json:"policyArn"` + Name string `json:"name"` + DefinitionHash string `json:"definitionHash"` + UpdatedAt string `json:"updatedAt,omitempty"` +} + +// --- dispatcher --- + +// serveRegistries routes the inference-profile, prompt-router, and +// automated-reasoning-policy control-plane surfaces. Split out of ServeHTTP to +// keep each dispatcher small. +func (h *Handler) serveRegistries(w http.ResponseWriter, r *http.Request, p string) { + switch { + case p == prefixInferenceProfiles || strings.HasPrefix(p, prefixInferenceProfiles+"/"): + h.serveInferenceProfiles(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixInferenceProfiles), "/")) + case p == prefixPromptRouters || strings.HasPrefix(p, prefixPromptRouters+"/"): + h.servePromptRouters(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixPromptRouters), "/")) + case p == prefixARPolicies || strings.HasPrefix(p, prefixARPolicies+"/"): + h.serveARPolicies(w, r, strings.TrimPrefix(strings.TrimPrefix(p, prefixARPolicies), "/")) + default: + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported path: "+p) + } +} + +// --- inference profile dispatch + operations --- + +// serveInferenceProfiles handles /inference-profiles[/{inferenceProfileIdentifier}]. +func (h *Handler) serveInferenceProfiles(w http.ResponseWriter, r *http.Request, id string) { + if id == "" { + switch r.Method { + case http.MethodPost: + h.createInferenceProfile(w, r) + case http.MethodGet: + h.listInferenceProfiles(w, r) + default: + methodNotAllowed(w) + } + + return + } + + switch r.Method { + case http.MethodGet: + h.getInferenceProfile(w, r, id) + case http.MethodDelete: + h.deleteInferenceProfile(w, r, id) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) createInferenceProfile(w http.ResponseWriter, r *http.Request) { + var in createInferenceProfileRequest + if !decodeJSON(w, r, &in) { + return + } + + p, err := h.bedrock.CreateInferenceProfile(r.Context(), bedrockdriver.InferenceProfileConfig{ + Name: in.InferenceProfileName, + ModelSourceCopyFrom: modelSourceCopyFrom(in.ModelSource), + ClientRequestToken: in.ClientRequestToken, + Description: in.Description, + Tags: tagsToMap(in.Tags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createInferenceProfileResponse{InferenceProfileARN: p.ARN, Status: p.Status}) +} + +func (h *Handler) getInferenceProfile(w http.ResponseWriter, r *http.Request, id string) { + p, err := h.bedrock.GetInferenceProfile(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toInferenceProfileJSON(p)) +} + +func (h *Handler) listInferenceProfiles(w http.ResponseWriter, r *http.Request) { + profiles, err := h.bedrock.ListInferenceProfiles(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]inferenceProfileJSON, 0, len(profiles)) + for i := range profiles { + out = append(out, toInferenceProfileJSON(&profiles[i])) + } + + writeJSON(w, listInferenceProfilesResponse{InferenceProfileSummaries: out}) +} + +func (h *Handler) deleteInferenceProfile(w http.ResponseWriter, r *http.Request, id string) { + if err := h.bedrock.DeleteInferenceProfile(r.Context(), id); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +// --- prompt router dispatch + operations --- + +// servePromptRouters handles /prompt-routers[/{promptRouterArn}]. The router ARN +// contains slashes, so it is the entire remainder of the path. +func (h *Handler) servePromptRouters(w http.ResponseWriter, r *http.Request, arn string) { + if arn == "" { + switch r.Method { + case http.MethodPost: + h.createPromptRouter(w, r) + case http.MethodGet: + h.listPromptRouters(w, r) + default: + methodNotAllowed(w) + } + + return + } + + switch r.Method { + case http.MethodGet: + h.getPromptRouter(w, r, arn) + case http.MethodDelete: + h.deletePromptRouter(w, r, arn) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) createPromptRouter(w http.ResponseWriter, r *http.Request) { + var in createPromptRouterRequest + if !decodeJSON(w, r, &in) { + return + } + + p, err := h.bedrock.CreatePromptRouter(r.Context(), bedrockdriver.PromptRouterConfig{ + Name: in.PromptRouterName, + Models: modelARNsOf(in.Models), + ResponseQualityDifference: responseQualityDifference(in.RoutingCriteria), + FallbackModelARN: fallbackModelARN(in.FallbackModel), + ClientRequestToken: in.ClientRequestToken, + Description: in.Description, + Tags: tagsToMap(in.Tags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createPromptRouterResponse{PromptRouterARN: p.ARN}) +} + +func (h *Handler) getPromptRouter(w http.ResponseWriter, r *http.Request, arn string) { + p, err := h.bedrock.GetPromptRouter(r.Context(), arn) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toPromptRouterJSON(p)) +} + +func (h *Handler) listPromptRouters(w http.ResponseWriter, r *http.Request) { + routers, err := h.bedrock.ListPromptRouters(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]promptRouterJSON, 0, len(routers)) + for i := range routers { + out = append(out, toPromptRouterJSON(&routers[i])) + } + + writeJSON(w, listPromptRoutersResponse{PromptRouterSummaries: out}) +} + +func (h *Handler) deletePromptRouter(w http.ResponseWriter, r *http.Request, arn string) { + if err := h.bedrock.DeletePromptRouter(r.Context(), arn); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +// --- automated reasoning policy dispatch + operations --- + +// serveARPolicies handles /automated-reasoning-policies[/{policyArn}]. The policy +// ARN contains slashes, so it is the entire remainder of the path. +func (h *Handler) serveARPolicies(w http.ResponseWriter, r *http.Request, arn string) { + if arn != "" { + switch r.Method { + case http.MethodGet: + h.getARPolicy(w, r, arn) + case http.MethodPatch: + h.updateARPolicy(w, r, arn) + case http.MethodDelete: + h.deleteARPolicy(w, r, arn) + default: + methodNotAllowed(w) + } + + return + } + + switch r.Method { + case http.MethodPost: + h.createARPolicy(w, r) + case http.MethodGet: + h.listARPolicies(w, r) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) createARPolicy(w http.ResponseWriter, r *http.Request) { + var in createARPolicyRequest + if !decodeJSON(w, r, &in) { + return + } + + p, err := h.bedrock.CreateAutomatedReasoningPolicy(r.Context(), bedrockdriver.AutomatedReasoningPolicyConfig{ + Name: in.Name, + ClientRequestToken: in.ClientRequestToken, + Description: in.Description, + KMSKeyID: in.KMSKeyID, + PolicyDefinition: []byte(in.PolicyDefinition), + Tags: tagsToMap(in.Tags), + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, createARPolicyResponse{ + PolicyARN: p.ARN, + Name: p.Name, + Version: p.Version, + DefinitionHash: p.DefinitionHash, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + Description: p.Description, + }) +} + +func (h *Handler) getARPolicy(w http.ResponseWriter, r *http.Request, arn string) { + p, err := h.bedrock.GetAutomatedReasoningPolicy(r.Context(), arn) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toARPolicyJSON(p)) +} + +func (h *Handler) listARPolicies(w http.ResponseWriter, r *http.Request) { + policies, err := h.bedrock.ListAutomatedReasoningPolicies(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]arPolicySummaryJSON, 0, len(policies)) + for i := range policies { + out = append(out, toARPolicySummaryJSON(&policies[i])) + } + + writeJSON(w, listARPoliciesResponse{AutomatedReasoningPolicySummaries: out}) +} + +func (h *Handler) updateARPolicy(w http.ResponseWriter, r *http.Request, arn string) { + var in updateARPolicyRequest + if !decodeJSON(w, r, &in) { + return + } + + p, err := h.bedrock.UpdateAutomatedReasoningPolicy(r.Context(), arn, bedrockdriver.AutomatedReasoningPolicyUpdate{ + PolicyDefinition: []byte(in.PolicyDefinition), + Description: in.Description, + Name: in.Name, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, updateARPolicyResponse{ + PolicyARN: p.ARN, + Name: p.Name, + DefinitionHash: p.DefinitionHash, + UpdatedAt: p.UpdatedAt, + }) +} + +func (h *Handler) deleteARPolicy(w http.ResponseWriter, r *http.Request, arn string) { + if err := h.bedrock.DeleteAutomatedReasoningPolicy(r.Context(), arn); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +// --- converters --- + +func modelSourceCopyFrom(in *inferenceProfileModelSourceJSON) string { + if in == nil { + return "" + } + + return in.CopyFrom +} + +func modelARNsOf(models []promptRouterTargetModelJSON) []string { + out := make([]string, 0, len(models)) + for _, m := range models { + out = append(out, m.ModelARN) + } + + return out +} + +func fallbackModelARN(in *promptRouterTargetModelJSON) string { + if in == nil { + return "" + } + + return in.ModelARN +} + +func responseQualityDifference(in *routingCriteriaJSON) *float64 { + if in == nil { + return nil + } + + return in.ResponseQualityDifference +} + +func targetModelsOf(arns []string) []promptRouterTargetModelJSON { + if len(arns) == 0 { + return nil + } + + out := make([]promptRouterTargetModelJSON, 0, len(arns)) + for _, a := range arns { + out = append(out, promptRouterTargetModelJSON{ModelARN: a}) + } + + return out +} + +func toInferenceProfileJSON(p *bedrockdriver.InferenceProfile) inferenceProfileJSON { + models := make([]inferenceProfileModelJSON, 0, len(p.Models)) + for _, a := range p.Models { + models = append(models, inferenceProfileModelJSON{ModelARN: a}) + } + + return inferenceProfileJSON{ + InferenceProfileARN: p.ARN, + InferenceProfileID: p.ID, + InferenceProfileName: p.Name, + Models: models, + Status: p.Status, + Type: p.Type, + Description: p.Description, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + } +} + +func toPromptRouterJSON(p *bedrockdriver.PromptRouter) promptRouterJSON { + out := promptRouterJSON{ + PromptRouterARN: p.ARN, + PromptRouterName: p.Name, + Models: targetModelsOf(p.Models), + Status: p.Status, + Type: p.Type, + Description: p.Description, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + } + if p.ResponseQualityDifference != nil { + out.RoutingCriteria = &routingCriteriaJSON{ResponseQualityDifference: p.ResponseQualityDifference} + } + + if p.FallbackModelARN != "" { + out.FallbackModel = &promptRouterTargetModelJSON{ModelARN: p.FallbackModelARN} + } + + return out +} + +func toARPolicyJSON(p *bedrockdriver.AutomatedReasoningPolicy) arPolicyJSON { + return arPolicyJSON{ + PolicyARN: p.ARN, + PolicyID: p.ID, + Name: p.Name, + Version: p.Version, + DefinitionHash: p.DefinitionHash, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + Description: p.Description, + KMSKeyARN: p.KMSKeyARN, + } +} + +func toARPolicySummaryJSON(p *bedrockdriver.AutomatedReasoningPolicy) arPolicySummaryJSON { + return arPolicySummaryJSON{ + PolicyARN: p.ARN, + PolicyID: p.ID, + Name: p.Name, + Version: p.Version, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + Description: p.Description, + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go b/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go new file mode 100644 index 00000000..f78be1ab --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go @@ -0,0 +1,275 @@ +package bedrock_test + +import ( + "context" + "errors" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsbedrock "github.com/aws/aws-sdk-go-v2/service/bedrock" + bedrocktypes "github.com/aws/aws-sdk-go-v2/service/bedrock/types" + awsruntime "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/document" + runtimetypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" + "github.com/aws/smithy-go" +) + +func TestSDKAsyncInvokeLifecycle(t *testing.T) { + client := newRuntimeClient(t) + ctx := context.Background() + + start, err := client.StartAsyncInvoke(ctx, &awsruntime.StartAsyncInvokeInput{ + ModelId: aws.String(claudeModel), + ModelInput: document.NewLazyDocument(map[string]any{"inputText": "hello"}), + OutputDataConfig: &runtimetypes.AsyncInvokeOutputDataConfigMemberS3OutputDataConfig{ + Value: runtimetypes.AsyncInvokeS3OutputDataConfig{S3Uri: aws.String("s3://bucket/out/")}, + }, + }) + if err != nil { + t.Fatalf("StartAsyncInvoke: %v", err) + } + + arn := aws.ToString(start.InvocationArn) + if arn == "" { + t.Fatal("expected an invocation ARN") + } + + got, err := client.GetAsyncInvoke(ctx, &awsruntime.GetAsyncInvokeInput{InvocationArn: aws.String(arn)}) + if err != nil { + t.Fatalf("GetAsyncInvoke: %v", err) + } + + if got.Status != runtimetypes.AsyncInvokeStatusCompleted { + t.Fatalf("got status %q, want Completed", got.Status) + } + + if aws.ToString(got.ModelArn) == "" || got.SubmitTime == nil { + t.Fatalf("expected modelArn + submitTime, got %+v", got) + } + + s3, ok := got.OutputDataConfig.(*runtimetypes.AsyncInvokeOutputDataConfigMemberS3OutputDataConfig) + if !ok || aws.ToString(s3.Value.S3Uri) != "s3://bucket/out/" { + t.Fatalf("unexpected output data config: %+v", got.OutputDataConfig) + } + + list, err := client.ListAsyncInvokes(ctx, &awsruntime.ListAsyncInvokesInput{}) + if err != nil { + t.Fatalf("ListAsyncInvokes: %v", err) + } + + if len(list.AsyncInvokeSummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.AsyncInvokeSummaries)) + } +} + +func TestSDKGetAsyncInvokeNotFound(t *testing.T) { + client := newRuntimeClient(t) + + _, err := client.GetAsyncInvoke(context.Background(), &awsruntime.GetAsyncInvokeInput{ + InvocationArn: aws.String("arn:aws:bedrock:us-east-1:123456789012:async-invoke/missing"), + }) + + var ae smithy.APIError + if !errors.As(err, &ae) { + t.Fatalf("expected API error, got %T: %v", err, err) + } + + if ae.ErrorCode() != "ResourceNotFoundException" { + t.Fatalf("got error code %q, want ResourceNotFoundException", ae.ErrorCode()) + } +} + +func TestSDKModelImportJobLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + create, err := client.CreateModelImportJob(ctx, &awsbedrock.CreateModelImportJobInput{ + JobName: aws.String("import-1"), + ImportedModelName: aws.String("my-imported"), + RoleArn: aws.String("arn:aws:iam::123456789012:role/bedrock"), + ModelDataSource: &bedrocktypes.ModelDataSourceMemberS3DataSource{ + Value: bedrocktypes.S3DataSource{S3Uri: aws.String("s3://bucket/model/")}, + }, + }) + if err != nil { + t.Fatalf("CreateModelImportJob: %v", err) + } + + if aws.ToString(create.JobArn) == "" { + t.Fatal("expected a job ARN") + } + + got, err := client.GetModelImportJob(ctx, &awsbedrock.GetModelImportJobInput{ + JobIdentifier: aws.String("import-1"), + }) + if err != nil { + t.Fatalf("GetModelImportJob: %v", err) + } + + if got.Status != bedrocktypes.ModelImportJobStatusCompleted { + t.Fatalf("got status %q, want Completed", got.Status) + } + + if aws.ToString(got.ImportedModelArn) == "" { + t.Fatalf("expected an imported model ARN, got %+v", got) + } + + if got.ModelDataSource == nil { + t.Fatal("expected modelDataSource to round-trip") + } + + list, err := client.ListModelImportJobs(ctx, &awsbedrock.ListModelImportJobsInput{}) + if err != nil { + t.Fatalf("ListModelImportJobs: %v", err) + } + + if len(list.ModelImportJobSummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.ModelImportJobSummaries)) + } + + _, err = client.GetModelImportJob(ctx, &awsbedrock.GetModelImportJobInput{JobIdentifier: aws.String("missing")}) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} + +func TestSDKModelCopyJobLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + src := "arn:aws:bedrock:us-east-1:123456789012:custom-model/src" + create, err := client.CreateModelCopyJob(ctx, &awsbedrock.CreateModelCopyJobInput{ + SourceModelArn: aws.String(src), + TargetModelName: aws.String("copy-target"), + }) + if err != nil { + t.Fatalf("CreateModelCopyJob: %v", err) + } + + arn := aws.ToString(create.JobArn) + if arn == "" { + t.Fatal("expected a job ARN") + } + + got, err := client.GetModelCopyJob(ctx, &awsbedrock.GetModelCopyJobInput{JobArn: aws.String(arn)}) + if err != nil { + t.Fatalf("GetModelCopyJob: %v", err) + } + + if got.Status != bedrocktypes.ModelCopyJobStatusCompleted { + t.Fatalf("got status %q, want Completed", got.Status) + } + + if aws.ToString(got.SourceModelArn) != src || aws.ToString(got.TargetModelArn) == "" { + t.Fatalf("unexpected copy job: %+v", got) + } + + if aws.ToString(got.SourceAccountId) != "123456789012" { + t.Fatalf("got source account %q", aws.ToString(got.SourceAccountId)) + } + + list, err := client.ListModelCopyJobs(ctx, &awsbedrock.ListModelCopyJobsInput{}) + if err != nil { + t.Fatalf("ListModelCopyJobs: %v", err) + } + + if len(list.ModelCopyJobSummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.ModelCopyJobSummaries)) + } + + _, err = client.GetModelCopyJob(ctx, &awsbedrock.GetModelCopyJobInput{ + JobArn: aws.String("arn:aws:bedrock:us-east-1:123456789012:model-copy-job/missing"), + }) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} + +func evaluationConfig() bedrocktypes.EvaluationConfig { + return &bedrocktypes.EvaluationConfigMemberAutomated{ + Value: bedrocktypes.AutomatedEvaluationConfig{ + DatasetMetricConfigs: []bedrocktypes.EvaluationDatasetMetricConfig{ + { + TaskType: bedrocktypes.EvaluationTaskTypeGeneration, + Dataset: &bedrocktypes.EvaluationDataset{Name: aws.String("ds")}, + MetricNames: []string{"Builtin.Accuracy"}, + }, + }, + }, + } +} + +func inferenceConfig() bedrocktypes.EvaluationInferenceConfig { + return &bedrocktypes.EvaluationInferenceConfigMemberModels{ + Value: []bedrocktypes.EvaluationModelConfig{ + &bedrocktypes.EvaluationModelConfigMemberBedrockModel{ + Value: bedrocktypes.EvaluationBedrockModel{ModelIdentifier: aws.String(claudeModel)}, + }, + }, + } +} + +func TestSDKEvaluationJobLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + create, err := client.CreateEvaluationJob(ctx, &awsbedrock.CreateEvaluationJobInput{ + JobName: aws.String("eval-1"), + RoleArn: aws.String("arn:aws:iam::123456789012:role/bedrock"), + EvaluationConfig: evaluationConfig(), + InferenceConfig: inferenceConfig(), + OutputDataConfig: &bedrocktypes.EvaluationOutputDataConfig{S3Uri: aws.String("s3://bucket/eval/")}, + }) + if err != nil { + t.Fatalf("CreateEvaluationJob: %v", err) + } + + if aws.ToString(create.JobArn) == "" { + t.Fatal("expected a job ARN") + } + + got, err := client.GetEvaluationJob(ctx, &awsbedrock.GetEvaluationJobInput{JobIdentifier: aws.String("eval-1")}) + if err != nil { + t.Fatalf("GetEvaluationJob: %v", err) + } + + if got.Status != bedrocktypes.EvaluationJobStatusCompleted { + t.Fatalf("got status %q, want Completed", got.Status) + } + + if got.JobType != bedrocktypes.EvaluationJobTypeAutomated { + t.Fatalf("got job type %q, want Automated", got.JobType) + } + + if _, ok := got.EvaluationConfig.(*bedrocktypes.EvaluationConfigMemberAutomated); !ok { + t.Fatalf("expected automated evaluation config to round-trip, got %T", got.EvaluationConfig) + } + + list, err := client.ListEvaluationJobs(ctx, &awsbedrock.ListEvaluationJobsInput{}) + if err != nil { + t.Fatalf("ListEvaluationJobs: %v", err) + } + + if len(list.JobSummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.JobSummaries)) + } + + if _, err = client.StopEvaluationJob(ctx, &awsbedrock.StopEvaluationJobInput{ + JobIdentifier: aws.String("eval-1"), + }); err != nil { + t.Fatalf("StopEvaluationJob: %v", err) + } + + stopped, err := client.GetEvaluationJob(ctx, &awsbedrock.GetEvaluationJobInput{JobIdentifier: aws.String("eval-1")}) + if err != nil { + t.Fatalf("GetEvaluationJob after stop: %v", err) + } + + if stopped.Status != bedrocktypes.EvaluationJobStatusStopped { + t.Fatalf("got status %q, want Stopped", stopped.Status) + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_guardrail_policies_test.go b/server/aws/bedrock/sdk_roundtrip_guardrail_policies_test.go new file mode 100644 index 00000000..be282da7 --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_guardrail_policies_test.go @@ -0,0 +1,152 @@ +package bedrock_test + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsbedrock "github.com/aws/aws-sdk-go-v2/service/bedrock" + bedrocktypes "github.com/aws/aws-sdk-go-v2/service/bedrock/types" +) + +// guardrailWithPolicies drives CreateGuardrail with a topic, content, PII, and +// contextual-grounding policy through the real SDK and returns the created id. +func guardrailWithPolicies(ctx context.Context, t *testing.T, client *awsbedrock.Client) string { + t.Helper() + + created, err := client.CreateGuardrail(ctx, &awsbedrock.CreateGuardrailInput{ + Name: aws.String("gr-policies"), + Description: aws.String("guardrail with policies"), + BlockedInputMessaging: aws.String("blocked input"), + BlockedOutputsMessaging: aws.String("blocked output"), + TopicPolicyConfig: &bedrocktypes.GuardrailTopicPolicyConfig{ + TopicsConfig: []bedrocktypes.GuardrailTopicConfig{{ + Name: aws.String("fiduciary-advice"), + Definition: aws.String("Providing personalized financial advice."), + Examples: []string{"Should I invest in this stock?"}, + Type: bedrocktypes.GuardrailTopicTypeDeny, + }}, + }, + ContentPolicyConfig: &bedrocktypes.GuardrailContentPolicyConfig{ + FiltersConfig: []bedrocktypes.GuardrailContentFilterConfig{{ + Type: bedrocktypes.GuardrailContentFilterTypeHate, + InputStrength: bedrocktypes.GuardrailFilterStrengthHigh, + OutputStrength: bedrocktypes.GuardrailFilterStrengthMedium, + }}, + }, + SensitiveInformationPolicyConfig: &bedrocktypes.GuardrailSensitiveInformationPolicyConfig{ + PiiEntitiesConfig: []bedrocktypes.GuardrailPiiEntityConfig{{ + Type: bedrocktypes.GuardrailPiiEntityTypeEmail, + Action: bedrocktypes.GuardrailSensitiveInformationActionAnonymize, + }}, + }, + ContextualGroundingPolicyConfig: &bedrocktypes.GuardrailContextualGroundingPolicyConfig{ + FiltersConfig: []bedrocktypes.GuardrailContextualGroundingFilterConfig{{ + Type: bedrocktypes.GuardrailContextualGroundingFilterTypeGrounding, + Threshold: aws.Float64(0.75), + Action: bedrocktypes.GuardrailContextualGroundingActionBlock, + }}, + }, + }) + if err != nil { + t.Fatalf("CreateGuardrail: %v", err) + } + + return aws.ToString(created.GuardrailId) +} + +func assertGuardrailPolicies(t *testing.T, got *awsbedrock.GetGuardrailOutput) { + t.Helper() + + if got.TopicPolicy == nil || len(got.TopicPolicy.Topics) != 1 || + aws.ToString(got.TopicPolicy.Topics[0].Name) != "fiduciary-advice" { + t.Fatalf("topic policy missing/wrong: %+v", got.TopicPolicy) + } + + if got.ContentPolicy == nil || len(got.ContentPolicy.Filters) != 1 || + got.ContentPolicy.Filters[0].Type != bedrocktypes.GuardrailContentFilterTypeHate { + t.Fatalf("content policy missing/wrong: %+v", got.ContentPolicy) + } + + if got.SensitiveInformationPolicy == nil || len(got.SensitiveInformationPolicy.PiiEntities) != 1 || + got.SensitiveInformationPolicy.PiiEntities[0].Action != bedrocktypes.GuardrailSensitiveInformationActionAnonymize { + t.Fatalf("sensitive-info policy missing/wrong: %+v", got.SensitiveInformationPolicy) + } + + if got.ContextualGroundingPolicy == nil || len(got.ContextualGroundingPolicy.Filters) != 1 || + aws.ToFloat64(got.ContextualGroundingPolicy.Filters[0].Threshold) != 0.75 { + t.Fatalf("contextual-grounding policy missing/wrong: %+v", got.ContextualGroundingPolicy) + } +} + +func TestSDKGuardrailPoliciesAndVersions(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + id := guardrailWithPolicies(ctx, t, client) + + // GetGuardrail (DRAFT) returns all policies with the response-shaped keys. + got, err := client.GetGuardrail(ctx, &awsbedrock.GetGuardrailInput{GuardrailIdentifier: aws.String(id)}) + if err != nil { + t.Fatalf("GetGuardrail: %v", err) + } + + assertGuardrailPolicies(t, got) + + // Snapshot the DRAFT into an immutable numbered version. + ver, err := client.CreateGuardrailVersion(ctx, &awsbedrock.CreateGuardrailVersionInput{ + GuardrailIdentifier: aws.String(id), + Description: aws.String("v1 snapshot"), + }) + if err != nil { + t.Fatalf("CreateGuardrailVersion: %v", err) + } + + if aws.ToString(ver.GuardrailId) == "" || aws.ToString(ver.Version) != "1" { + t.Fatalf("expected guardrail id + version 1, got %+v", ver) + } + + // GetGuardrail at the numbered version returns the snapshot with policies. + snap, err := client.GetGuardrail(ctx, &awsbedrock.GetGuardrailInput{ + GuardrailIdentifier: aws.String(id), + GuardrailVersion: aws.String("1"), + }) + if err != nil { + t.Fatalf("GetGuardrail(version): %v", err) + } + + if aws.ToString(snap.Version) != "1" { + t.Fatalf("got version %q, want 1", aws.ToString(snap.Version)) + } + + assertGuardrailPolicies(t, snap) + + // Scoped list shows DRAFT plus the numbered version. + list, err := client.ListGuardrails(ctx, &awsbedrock.ListGuardrailsInput{GuardrailIdentifier: aws.String(id)}) + if err != nil { + t.Fatalf("ListGuardrails: %v", err) + } + + if len(list.Guardrails) != 2 { + t.Fatalf("got %d guardrail summaries, want 2 (DRAFT + v1)", len(list.Guardrails)) + } + + // Deleting the specific version leaves the DRAFT resolvable. + if _, err = client.DeleteGuardrail(ctx, &awsbedrock.DeleteGuardrailInput{ + GuardrailIdentifier: aws.String(id), + GuardrailVersion: aws.String("1"), + }); err != nil { + t.Fatalf("DeleteGuardrail(version): %v", err) + } + + if _, err = client.GetGuardrail(ctx, &awsbedrock.GetGuardrailInput{ + GuardrailIdentifier: aws.String(id), + GuardrailVersion: aws.String("1"), + }); err == nil { + t.Fatal("expected error getting deleted version") + } + + if _, err = client.GetGuardrail(ctx, &awsbedrock.GetGuardrailInput{GuardrailIdentifier: aws.String(id)}); err != nil { + t.Fatalf("DRAFT should survive version delete: %v", err) + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_marketplace_test.go b/server/aws/bedrock/sdk_roundtrip_marketplace_test.go new file mode 100644 index 00000000..bfbcaf06 --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_marketplace_test.go @@ -0,0 +1,189 @@ +package bedrock_test + +import ( + "context" + "errors" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsbedrock "github.com/aws/aws-sdk-go-v2/service/bedrock" + bedrocktypes "github.com/aws/aws-sdk-go-v2/service/bedrock/types" +) + +func TestSDKMarketplaceModelEndpointLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + source := "arn:aws:sagemaker:us-east-1:aws:hub-content/SageMakerPublicHub/Model/model-1" + cfg := &bedrocktypes.EndpointConfigMemberSageMaker{Value: bedrocktypes.SageMakerEndpoint{ + ExecutionRole: aws.String("arn:aws:iam::123456789012:role/sagemaker"), + InitialInstanceCount: aws.Int32(1), + InstanceType: aws.String("ml.m5.large"), + }} + + create, err := client.CreateMarketplaceModelEndpoint(ctx, &awsbedrock.CreateMarketplaceModelEndpointInput{ + EndpointConfig: cfg, + EndpointName: aws.String("endpoint-1"), + ModelSourceIdentifier: aws.String(source), + AcceptEula: true, + }) + if err != nil { + t.Fatalf("CreateMarketplaceModelEndpoint: %v", err) + } + + endpoint := create.MarketplaceModelEndpoint + arn := aws.ToString(endpoint.EndpointArn) + if arn == "" { + t.Fatal("expected an endpoint ARN") + } + + if endpoint.Status != bedrocktypes.StatusRegistered { + t.Fatalf("got status %q, want REGISTERED", endpoint.Status) + } + + if aws.ToString(endpoint.EndpointStatus) != "InService" { + t.Fatalf("got endpointStatus %q, want InService", aws.ToString(endpoint.EndpointStatus)) + } + + got, err := client.GetMarketplaceModelEndpoint(ctx, &awsbedrock.GetMarketplaceModelEndpointInput{ + EndpointArn: aws.String(arn), + }) + if err != nil { + t.Fatalf("GetMarketplaceModelEndpoint: %v", err) + } + + sm, ok := got.MarketplaceModelEndpoint.EndpointConfig.(*bedrocktypes.EndpointConfigMemberSageMaker) + if !ok || aws.ToString(sm.Value.InstanceType) != "ml.m5.large" { + t.Fatalf("unexpected endpoint config: %+v", got.MarketplaceModelEndpoint.EndpointConfig) + } + + list, err := client.ListMarketplaceModelEndpoints(ctx, &awsbedrock.ListMarketplaceModelEndpointsInput{}) + if err != nil { + t.Fatalf("ListMarketplaceModelEndpoints: %v", err) + } + + if len(list.MarketplaceModelEndpoints) != 1 || aws.ToString(list.MarketplaceModelEndpoints[0].EndpointArn) != arn { + t.Fatalf("unexpected list result: %+v", list.MarketplaceModelEndpoints) + } + + upd := &bedrocktypes.EndpointConfigMemberSageMaker{Value: bedrocktypes.SageMakerEndpoint{ + ExecutionRole: aws.String("arn:aws:iam::123456789012:role/sagemaker"), + InitialInstanceCount: aws.Int32(2), + InstanceType: aws.String("ml.m5.xlarge"), + }} + + updated, err := client.UpdateMarketplaceModelEndpoint(ctx, &awsbedrock.UpdateMarketplaceModelEndpointInput{ + EndpointArn: aws.String(arn), + EndpointConfig: upd, + }) + if err != nil { + t.Fatalf("UpdateMarketplaceModelEndpoint: %v", err) + } + + usm, ok := updated.MarketplaceModelEndpoint.EndpointConfig.(*bedrocktypes.EndpointConfigMemberSageMaker) + if !ok || aws.ToString(usm.Value.InstanceType) != "ml.m5.xlarge" { + t.Fatalf("unexpected updated config: %+v", updated.MarketplaceModelEndpoint.EndpointConfig) + } + + reg, err := client.RegisterMarketplaceModelEndpoint(ctx, &awsbedrock.RegisterMarketplaceModelEndpointInput{ + EndpointIdentifier: aws.String(arn), + ModelSourceIdentifier: aws.String(source), + }) + if err != nil { + t.Fatalf("RegisterMarketplaceModelEndpoint: %v", err) + } + + if reg.MarketplaceModelEndpoint.Status != bedrocktypes.StatusRegistered { + t.Fatalf("got status %q after register, want REGISTERED", reg.MarketplaceModelEndpoint.Status) + } + + if _, err = client.DeregisterMarketplaceModelEndpoint(ctx, &awsbedrock.DeregisterMarketplaceModelEndpointInput{ + EndpointArn: aws.String(arn), + }); err != nil { + t.Fatalf("DeregisterMarketplaceModelEndpoint: %v", err) + } + + if _, err = client.DeleteMarketplaceModelEndpoint(ctx, &awsbedrock.DeleteMarketplaceModelEndpointInput{ + EndpointArn: aws.String(arn), + }); err != nil { + t.Fatalf("DeleteMarketplaceModelEndpoint: %v", err) + } + + _, err = client.GetMarketplaceModelEndpoint(ctx, &awsbedrock.GetMarketplaceModelEndpointInput{ + EndpointArn: aws.String(arn), + }) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} + +func TestSDKFoundationModelAgreementLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + offers, err := client.ListFoundationModelAgreementOffers(ctx, &awsbedrock.ListFoundationModelAgreementOffersInput{ + ModelId: aws.String(claudeModel), + OfferType: bedrocktypes.OfferTypeAll, + }) + if err != nil { + t.Fatalf("ListFoundationModelAgreementOffers: %v", err) + } + + if len(offers.Offers) != 1 || aws.ToString(offers.Offers[0].OfferToken) == "" { + t.Fatalf("unexpected offers: %+v", offers.Offers) + } + + create, err := client.CreateFoundationModelAgreement(ctx, &awsbedrock.CreateFoundationModelAgreementInput{ + ModelId: aws.String(claudeModel), + OfferToken: offers.Offers[0].OfferToken, + }) + if err != nil { + t.Fatalf("CreateFoundationModelAgreement: %v", err) + } + + if aws.ToString(create.ModelId) != claudeModel { + t.Fatalf("got modelId %q, want %q", aws.ToString(create.ModelId), claudeModel) + } + + avail, err := client.GetFoundationModelAvailability(ctx, &awsbedrock.GetFoundationModelAvailabilityInput{ + ModelId: aws.String(claudeModel), + }) + if err != nil { + t.Fatalf("GetFoundationModelAvailability: %v", err) + } + + if avail.AuthorizationStatus != bedrocktypes.AuthorizationStatusAuthorized { + t.Fatalf("got authorization %q, want AUTHORIZED", avail.AuthorizationStatus) + } + + if avail.AgreementAvailability == nil || avail.AgreementAvailability.Status != bedrocktypes.AgreementStatusAvailable { + t.Fatalf("unexpected agreement availability: %+v", avail.AgreementAvailability) + } + + if avail.EntitlementAvailability != bedrocktypes.EntitlementAvailabilityAvailable { + t.Fatalf("got entitlement %q, want AVAILABLE", avail.EntitlementAvailability) + } + + if avail.RegionAvailability != bedrocktypes.RegionAvailabilityAvailable { + t.Fatalf("got region %q, want AVAILABLE", avail.RegionAvailability) + } + + if _, err = client.DeleteFoundationModelAgreement(ctx, &awsbedrock.DeleteFoundationModelAgreementInput{ + ModelId: aws.String(claudeModel), + }); err != nil { + t.Fatalf("DeleteFoundationModelAgreement: %v", err) + } + + after, err := client.GetFoundationModelAvailability(ctx, &awsbedrock.GetFoundationModelAvailabilityInput{ + ModelId: aws.String(claudeModel), + }) + if err != nil { + t.Fatalf("GetFoundationModelAvailability after delete: %v", err) + } + + if after.AuthorizationStatus != bedrocktypes.AuthorizationStatusNotAuthorized { + t.Fatalf("got authorization %q, want NOT_AUTHORIZED", after.AuthorizationStatus) + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_registries_test.go b/server/aws/bedrock/sdk_roundtrip_registries_test.go new file mode 100644 index 00000000..a73d4d5f --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_registries_test.go @@ -0,0 +1,213 @@ +package bedrock_test + +import ( + "context" + "errors" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsbedrock "github.com/aws/aws-sdk-go-v2/service/bedrock" + bedrocktypes "github.com/aws/aws-sdk-go-v2/service/bedrock/types" +) + +func TestSDKInferenceProfileLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + src := "arn:aws:bedrock:us-east-1::foundation-model/" + claudeModel + create, err := client.CreateInferenceProfile(ctx, &awsbedrock.CreateInferenceProfileInput{ + InferenceProfileName: aws.String("profile-1"), + ModelSource: &bedrocktypes.InferenceProfileModelSourceMemberCopyFrom{Value: src}, + Description: aws.String("test profile"), + }) + if err != nil { + t.Fatalf("CreateInferenceProfile: %v", err) + } + + arn := aws.ToString(create.InferenceProfileArn) + if arn == "" { + t.Fatal("expected an inference profile ARN") + } + + if create.Status != bedrocktypes.InferenceProfileStatusActive { + t.Fatalf("got status %q, want ACTIVE", create.Status) + } + + got, err := client.GetInferenceProfile(ctx, &awsbedrock.GetInferenceProfileInput{ + InferenceProfileIdentifier: aws.String(arn), + }) + if err != nil { + t.Fatalf("GetInferenceProfile: %v", err) + } + + if got.Type != bedrocktypes.InferenceProfileTypeApplication { + t.Fatalf("got type %q, want APPLICATION", got.Type) + } + + if len(got.Models) != 1 || aws.ToString(got.Models[0].ModelArn) != src { + t.Fatalf("unexpected models: %+v", got.Models) + } + + list, err := client.ListInferenceProfiles(ctx, &awsbedrock.ListInferenceProfilesInput{}) + if err != nil { + t.Fatalf("ListInferenceProfiles: %v", err) + } + + if len(list.InferenceProfileSummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.InferenceProfileSummaries)) + } + + if _, err = client.DeleteInferenceProfile(ctx, &awsbedrock.DeleteInferenceProfileInput{ + InferenceProfileIdentifier: aws.String(arn), + }); err != nil { + t.Fatalf("DeleteInferenceProfile: %v", err) + } + + _, err = client.GetInferenceProfile(ctx, &awsbedrock.GetInferenceProfileInput{ + InferenceProfileIdentifier: aws.String(arn), + }) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} + +func TestSDKPromptRouterLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + diff := 0.5 + create, err := client.CreatePromptRouter(ctx, &awsbedrock.CreatePromptRouterInput{ + PromptRouterName: aws.String("router-1"), + Models: []bedrocktypes.PromptRouterTargetModel{ + {ModelArn: aws.String("arn:aws:bedrock:us-east-1::foundation-model/" + claudeModel)}, + }, + RoutingCriteria: &bedrocktypes.RoutingCriteria{ResponseQualityDifference: aws.Float64(diff)}, + FallbackModel: &bedrocktypes.PromptRouterTargetModel{ModelArn: aws.String("arn:model/fallback")}, + Description: aws.String("test router"), + }) + if err != nil { + t.Fatalf("CreatePromptRouter: %v", err) + } + + arn := aws.ToString(create.PromptRouterArn) + if arn == "" { + t.Fatal("expected a prompt router ARN") + } + + got, err := client.GetPromptRouter(ctx, &awsbedrock.GetPromptRouterInput{PromptRouterArn: aws.String(arn)}) + if err != nil { + t.Fatalf("GetPromptRouter: %v", err) + } + + if got.Status != bedrocktypes.PromptRouterStatusAvailable { + t.Fatalf("got status %q, want AVAILABLE", got.Status) + } + + if got.Type != bedrocktypes.PromptRouterTypeCustom { + t.Fatalf("got type %q, want custom", got.Type) + } + + if got.RoutingCriteria == nil || aws.ToFloat64(got.RoutingCriteria.ResponseQualityDifference) != diff { + t.Fatalf("unexpected routing criteria: %+v", got.RoutingCriteria) + } + + if got.FallbackModel == nil || aws.ToString(got.FallbackModel.ModelArn) != "arn:model/fallback" { + t.Fatalf("unexpected fallback model: %+v", got.FallbackModel) + } + + list, err := client.ListPromptRouters(ctx, &awsbedrock.ListPromptRoutersInput{}) + if err != nil { + t.Fatalf("ListPromptRouters: %v", err) + } + + if len(list.PromptRouterSummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.PromptRouterSummaries)) + } + + if _, err = client.DeletePromptRouter(ctx, &awsbedrock.DeletePromptRouterInput{ + PromptRouterArn: aws.String(arn), + }); err != nil { + t.Fatalf("DeletePromptRouter: %v", err) + } + + _, err = client.GetPromptRouter(ctx, &awsbedrock.GetPromptRouterInput{PromptRouterArn: aws.String(arn)}) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} + +func TestSDKAutomatedReasoningPolicyLifecycle(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + create, err := client.CreateAutomatedReasoningPolicy(ctx, &awsbedrock.CreateAutomatedReasoningPolicyInput{ + Name: aws.String("policy-1"), + Description: aws.String("test policy"), + PolicyDefinition: &bedrocktypes.AutomatedReasoningPolicyDefinition{Version: aws.String("1.0")}, + }) + if err != nil { + t.Fatalf("CreateAutomatedReasoningPolicy: %v", err) + } + + arn := aws.ToString(create.PolicyArn) + if arn == "" || aws.ToString(create.DefinitionHash) == "" { + t.Fatalf("expected policy ARN + definition hash, got %+v", create) + } + + if aws.ToString(create.Version) != "DRAFT" { + t.Fatalf("got version %q, want DRAFT", aws.ToString(create.Version)) + } + + got, err := client.GetAutomatedReasoningPolicy(ctx, &awsbedrock.GetAutomatedReasoningPolicyInput{ + PolicyArn: aws.String(arn), + }) + if err != nil { + t.Fatalf("GetAutomatedReasoningPolicy: %v", err) + } + + if aws.ToString(got.Name) != "policy-1" || aws.ToString(got.PolicyId) == "" { + t.Fatalf("unexpected policy: %+v", got) + } + + list, err := client.ListAutomatedReasoningPolicies(ctx, &awsbedrock.ListAutomatedReasoningPoliciesInput{}) + if err != nil { + t.Fatalf("ListAutomatedReasoningPolicies: %v", err) + } + + if len(list.AutomatedReasoningPolicySummaries) != 1 { + t.Fatalf("got %d summaries, want 1", len(list.AutomatedReasoningPolicySummaries)) + } + + updated, err := client.UpdateAutomatedReasoningPolicy(ctx, &awsbedrock.UpdateAutomatedReasoningPolicyInput{ + PolicyArn: aws.String(arn), + Name: aws.String("policy-1-renamed"), + Description: aws.String("updated"), + PolicyDefinition: &bedrocktypes.AutomatedReasoningPolicyDefinition{Version: aws.String("2.0")}, + }) + if err != nil { + t.Fatalf("UpdateAutomatedReasoningPolicy: %v", err) + } + + if aws.ToString(updated.Name) != "policy-1-renamed" || aws.ToString(updated.DefinitionHash) == "" { + t.Fatalf("unexpected update result: %+v", updated) + } + + if _, err = client.DeleteAutomatedReasoningPolicy(ctx, &awsbedrock.DeleteAutomatedReasoningPolicyInput{ + PolicyArn: aws.String(arn), + }); err != nil { + t.Fatalf("DeleteAutomatedReasoningPolicy: %v", err) + } + + _, err = client.GetAutomatedReasoningPolicy(ctx, &awsbedrock.GetAutomatedReasoningPolicyInput{ + PolicyArn: aws.String(arn), + }) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_runtime_test.go b/server/aws/bedrock/sdk_roundtrip_runtime_test.go new file mode 100644 index 00000000..119e771d --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_runtime_test.go @@ -0,0 +1,106 @@ +package bedrock_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsbedrock "github.com/aws/aws-sdk-go-v2/service/bedrock" + awsruntime "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + runtimetypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" +) + +func TestSDKCountTokensConverse(t *testing.T) { + client := newRuntimeClient(t) + + out, err := client.CountTokens(context.Background(), &awsruntime.CountTokensInput{ + ModelId: aws.String(claudeModel), + Input: &runtimetypes.CountTokensInputMemberConverse{ + Value: runtimetypes.ConverseTokensRequest{ + System: []runtimetypes.SystemContentBlock{ + &runtimetypes.SystemContentBlockMemberText{Value: "Be concise."}, + }, + Messages: []runtimetypes.Message{ + { + Role: runtimetypes.ConversationRoleUser, + Content: []runtimetypes.ContentBlock{&runtimetypes.ContentBlockMemberText{Value: "What is Bedrock?"}}, + }, + }, + }, + }, + }) + if err != nil { + t.Fatalf("CountTokens: %v", err) + } + + if aws.ToInt32(out.InputTokens) <= 0 { + t.Fatalf("expected a positive input token count, got %d", aws.ToInt32(out.InputTokens)) + } +} + +// TestSDKCountTokensInvokeModel exercises the invokeModel union member, whose +// body the SDK serializes as a base64 blob — the server must decode it back to +// the model-native payload before counting tokens. +func TestSDKCountTokensInvokeModel(t *testing.T) { + client := newRuntimeClient(t) + + body, _ := json.Marshal(map[string]any{ + "anthropic_version": "bedrock-2023-05-31", + "messages": []map[string]any{{"role": "user", "content": "how many tokens is this prompt"}}, + }) + + out, err := client.CountTokens(context.Background(), &awsruntime.CountTokensInput{ + ModelId: aws.String(claudeModel), + Input: &runtimetypes.CountTokensInputMemberInvokeModel{Value: runtimetypes.InvokeModelTokensRequest{Body: body}}, + }) + if err != nil { + t.Fatalf("CountTokens (invokeModel): %v", err) + } + + if aws.ToInt32(out.InputTokens) <= 0 { + t.Fatalf("expected a positive input token count on the invokeModel path, got %d", aws.ToInt32(out.InputTokens)) + } +} + +func TestSDKApplyGuardrail(t *testing.T) { + endpoint := newServer(t) + cfg := testConfig(t) + + control := awsbedrock.NewFromConfig(cfg, func(o *awsbedrock.Options) { + o.BaseEndpoint = aws.String(endpoint) + }) + runtime := awsruntime.NewFromConfig(cfg, func(o *awsruntime.Options) { + o.BaseEndpoint = aws.String(endpoint) + }) + + ctx := context.Background() + + created, err := control.CreateGuardrail(ctx, &awsbedrock.CreateGuardrailInput{ + Name: aws.String("gr-apply"), + Description: aws.String("apply guardrail test"), + BlockedInputMessaging: aws.String("blocked input"), + BlockedOutputsMessaging: aws.String("blocked output"), + }) + if err != nil { + t.Fatalf("CreateGuardrail: %v", err) + } + + out, err := runtime.ApplyGuardrail(ctx, &awsruntime.ApplyGuardrailInput{ + GuardrailIdentifier: created.GuardrailId, + GuardrailVersion: aws.String("DRAFT"), + Source: runtimetypes.GuardrailContentSourceInput, + Content: []runtimetypes.GuardrailContentBlock{ + &runtimetypes.GuardrailContentBlockMemberText{ + Value: runtimetypes.GuardrailTextBlock{Text: aws.String("hello guardrail")}, + }, + }, + }) + if err != nil { + t.Fatalf("ApplyGuardrail: %v", err) + } + + if out.Action != runtimetypes.GuardrailActionNone { + t.Fatalf("got action %q, want NONE", out.Action) + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_streaming_test.go b/server/aws/bedrock/sdk_roundtrip_streaming_test.go new file mode 100644 index 00000000..fe5e5bcd --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_streaming_test.go @@ -0,0 +1,118 @@ +package bedrock_test + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsruntime "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + runtimetypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" +) + +func TestSDKConverseStream(t *testing.T) { + client := newRuntimeClient(t) + + out, err := client.ConverseStream(context.Background(), &awsruntime.ConverseStreamInput{ + ModelId: aws.String(claudeModel), + Messages: []runtimetypes.Message{ + { + Role: runtimetypes.ConversationRoleUser, + Content: []runtimetypes.ContentBlock{&runtimetypes.ContentBlockMemberText{Value: "Stream me a reply."}}, + }, + }, + }) + if err != nil { + t.Fatalf("ConverseStream: %v", err) + } + + stream := out.GetStream() + defer stream.Close() + + var ( + text strings.Builder + sawStart bool + sawStop bool + sawMetadata bool + inTokens int32 + ) + + for ev := range stream.Events() { + switch e := ev.(type) { + case *runtimetypes.ConverseStreamOutputMemberMessageStart: + sawStart = true + case *runtimetypes.ConverseStreamOutputMemberContentBlockDelta: + if d, ok := e.Value.Delta.(*runtimetypes.ContentBlockDeltaMemberText); ok { + text.WriteString(d.Value) + } + case *runtimetypes.ConverseStreamOutputMemberMessageStop: + sawStop = true + case *runtimetypes.ConverseStreamOutputMemberMetadata: + sawMetadata = true + if e.Value.Usage != nil { + inTokens = aws.ToInt32(e.Value.Usage.InputTokens) + } + } + } + + if err := stream.Err(); err != nil { + t.Fatalf("stream error: %v", err) + } + + if !sawStart || !sawStop || !sawMetadata { + t.Fatalf("missing lifecycle events: start=%v stop=%v metadata=%v", sawStart, sawStop, sawMetadata) + } + + if text.Len() == 0 { + t.Fatal("expected non-empty streamed assistant text") + } + + if inTokens <= 0 { + t.Fatalf("expected positive input token usage, got %d", inTokens) + } +} + +func TestSDKInvokeModelWithResponseStream(t *testing.T) { + client := newRuntimeClient(t) + + body, _ := json.Marshal(map[string]any{ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 32, + "messages": []map[string]any{{"role": "user", "content": "Hi there"}}, + }) + + out, err := client.InvokeModelWithResponseStream(context.Background(), &awsruntime.InvokeModelWithResponseStreamInput{ + ModelId: aws.String(claudeModel), + ContentType: aws.String("application/json"), + Body: body, + }) + if err != nil { + t.Fatalf("InvokeModelWithResponseStream: %v", err) + } + + stream := out.GetStream() + defer stream.Close() + + var payload []byte + + for ev := range stream.Events() { + if c, ok := ev.(*runtimetypes.ResponseStreamMemberChunk); ok { + payload = append(payload, c.Value.Bytes...) + } + } + + if err := stream.Err(); err != nil { + t.Fatalf("stream error: %v", err) + } + + if len(payload) == 0 { + t.Fatal("expected a non-empty chunk payload") + } + + // The chunk carries model-native JSON; confirm it parses. + var probe map[string]any + if err := json.Unmarshal(payload, &probe); err != nil { + t.Fatalf("chunk bytes are not valid JSON: %v (%s)", err, string(payload)) + } +} diff --git a/server/aws/bedrock/sdk_roundtrip_tags_test.go b/server/aws/bedrock/sdk_roundtrip_tags_test.go new file mode 100644 index 00000000..becafe90 --- /dev/null +++ b/server/aws/bedrock/sdk_roundtrip_tags_test.go @@ -0,0 +1,86 @@ +package bedrock_test + +import ( + "context" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsbedrock "github.com/aws/aws-sdk-go-v2/service/bedrock" + bedrocktypes "github.com/aws/aws-sdk-go-v2/service/bedrock/types" +) + +func TestSDKTagResourceRoundtrip(t *testing.T) { + client := newControlClient(t) + ctx := context.Background() + + created, err := client.CreateGuardrail(ctx, &awsbedrock.CreateGuardrailInput{ + Name: aws.String("gr-tags"), + BlockedInputMessaging: aws.String("blocked input"), + BlockedOutputsMessaging: aws.String("blocked output"), + }) + if err != nil { + t.Fatalf("CreateGuardrail: %v", err) + } + + arn := aws.ToString(created.GuardrailArn) + if arn == "" { + t.Fatal("expected a guardrail ARN") + } + + if _, err = client.TagResource(ctx, &awsbedrock.TagResourceInput{ + ResourceARN: aws.String(arn), + Tags: []bedrocktypes.Tag{ + {Key: aws.String("env"), Value: aws.String("prod")}, + {Key: aws.String("team"), Value: aws.String("ml")}, + }, + }); err != nil { + t.Fatalf("TagResource: %v", err) + } + + list, err := client.ListTagsForResource(ctx, &awsbedrock.ListTagsForResourceInput{ + ResourceARN: aws.String(arn), + }) + if err != nil { + t.Fatalf("ListTagsForResource: %v", err) + } + + if got := tagValue(list.Tags, "env"); got != "prod" { + t.Fatalf("env tag = %q, want prod (tags: %+v)", got, list.Tags) + } + + if got := tagValue(list.Tags, "team"); got != "ml" { + t.Fatalf("team tag = %q, want ml (tags: %+v)", got, list.Tags) + } + + if _, err = client.UntagResource(ctx, &awsbedrock.UntagResourceInput{ + ResourceARN: aws.String(arn), + TagKeys: []string{"env"}, + }); err != nil { + t.Fatalf("UntagResource: %v", err) + } + + after, err := client.ListTagsForResource(ctx, &awsbedrock.ListTagsForResourceInput{ + ResourceARN: aws.String(arn), + }) + if err != nil { + t.Fatalf("ListTagsForResource after untag: %v", err) + } + + if got := tagValue(after.Tags, "env"); got != "" { + t.Fatalf("expected env tag to be removed, still got %q", got) + } + + if got := tagValue(after.Tags, "team"); got != "ml" { + t.Fatalf("team tag = %q, want ml after untag", got) + } +} + +func tagValue(tags []bedrocktypes.Tag, key string) string { + for _, tag := range tags { + if aws.ToString(tag.Key) == key { + return aws.ToString(tag.Value) + } + } + + return "" +} diff --git a/server/aws/bedrock/streaming.go b/server/aws/bedrock/streaming.go new file mode 100644 index 00000000..d5758dab --- /dev/null +++ b/server/aws/bedrock/streaming.go @@ -0,0 +1,157 @@ +package bedrock + +import ( + "encoding/base64" + "encoding/json" + "io" + "net/http" + "strings" + + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +const ( + contentTypeEventStream = "application/vnd.amazon.eventstream" + keyContentBlockIndex = "contentBlockIndex" + // minSplitLen is the shortest text worth splitting across two deltas. + minSplitLen = 2 +) + +// eventWriter frames Bedrock runtime streaming responses as +// application/vnd.amazon.eventstream events, flushing each event as it is +// written so SDK clients observe an incremental stream. +type eventWriter struct { + w http.ResponseWriter + enc *eventstream.Encoder + flusher http.Flusher +} + +// newEventWriter sets the eventstream content type, writes a 200 status, and +// returns a writer ready to emit events. Call only after the driver result is +// known so errors can still be reported via writeErr. +func newEventWriter(w http.ResponseWriter) *eventWriter { + w.Header().Set("Content-Type", contentTypeEventStream) + w.WriteHeader(http.StatusOK) + + flusher, _ := w.(http.Flusher) + + return &eventWriter{w: w, enc: eventstream.NewEncoder(), flusher: flusher} +} + +// event encodes one JSON event of the given type and flushes it. +func (e *eventWriter) event(eventType string, payload []byte) { + var h eventstream.Headers + + h.Set(eventstreamapi.MessageTypeHeader, eventstream.StringValue(eventstreamapi.EventMessageType)) + h.Set(eventstreamapi.EventTypeHeader, eventstream.StringValue(eventType)) + h.Set(eventstreamapi.ContentTypeHeader, eventstream.StringValue(contentTypeJSON)) + + _ = e.enc.Encode(e.w, eventstream.Message{Headers: h, Payload: payload}) + + if e.flusher != nil { + e.flusher.Flush() + } +} + +// converseStream handles POST /model/{modelId}/converse-stream. It reuses the +// Converse driver call and frames the full result as a sequence of streaming +// events (messageStart, contentBlockDelta(s), contentBlockStop, messageStop, +// metadata). +func (h *Handler) converseStream(w http.ResponseWriter, r *http.Request, modelID string) { + var in converseRequest + if !decodeJSON(w, r, &in) { + return + } + + out, err := h.bedrock.Converse(r.Context(), toConverseInput(modelID, &in)) + if err != nil { + writeErr(w, err) + + return + } + + ev := newEventWriter(w) + ev.event("messageStart", mustJSON(map[string]string{"role": out.Message.Role})) + + for _, chunk := range chunkText(strings.Join(out.Message.Text, "")) { + ev.event("contentBlockDelta", mustJSON(map[string]any{ + keyContentBlockIndex: 0, + "delta": map[string]string{"text": chunk}, + })) + } + + ev.event("contentBlockStop", mustJSON(map[string]any{keyContentBlockIndex: 0})) + ev.event("messageStop", mustJSON(map[string]string{"stopReason": out.StopReason})) + ev.event("metadata", metadataPayload(out)) +} + +// invokeModelStream handles POST /model/{modelId}/invoke-with-response-stream. +// It reuses the InvokeModel driver call and emits the model-native response as +// a single base64-encoded chunk event. +func (h *Handler) invokeModelStream(w http.ResponseWriter, r *http.Request, modelID string) { + r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) + + body, err := io.ReadAll(r.Body) + if err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", "read body: "+err.Error()) + + return + } + + res, err := h.bedrock.InvokeModel(r.Context(), bedrockdriver.InvokeModelInput{ + ModelID: modelID, + ContentType: r.Header.Get("Content-Type"), + Accept: r.Header.Get("Accept"), + Body: body, + }) + if err != nil { + writeErr(w, err) + + return + } + + ev := newEventWriter(w) + ev.event("chunk", mustJSON(map[string]string{ + "bytes": base64.StdEncoding.EncodeToString(res.Body), + })) +} + +// metadataPayload builds the Converse metadata event carrying token usage and +// latency metrics. +func metadataPayload(out *bedrockdriver.ConverseOutput) []byte { + return mustJSON(map[string]any{ + "usage": map[string]int{ + "inputTokens": out.InputTokens, + "outputTokens": out.OutputTokens, + "totalTokens": out.TotalTokens, + }, + "metrics": map[string]int{"latencyMs": out.LatencyMs}, + }) +} + +// chunkText splits s into up to two contentBlockDelta chunks so the emulated +// stream delivers more than one delta when the text is long enough. It always +// returns at least one chunk. +func chunkText(s string) []string { + if len(s) < minSplitLen { + return []string{s} + } + + mid := len(s) / 2 + + return []string{s[:mid], s[mid:]} +} + +// mustJSON marshals v to JSON, returning an empty object on the impossible +// error path for these fixed, marshalable payload shapes. +func mustJSON(v any) []byte { + b, err := json.Marshal(v) + if err != nil { + return []byte("{}") + } + + return b +} diff --git a/server/aws/bedrock/tags.go b/server/aws/bedrock/tags.go new file mode 100644 index 00000000..b417edcd --- /dev/null +++ b/server/aws/bedrock/tags.go @@ -0,0 +1,77 @@ +package bedrock + +import ( + "net/http" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// tagResource handles POST /tagResource, associating tags with a resource ARN. +func (h *Handler) tagResource(w http.ResponseWriter, r *http.Request) { + var in tagResourceRequest + if !decodeJSON(w, r, &in) { + return + } + + if err := h.bedrock.TagResource(r.Context(), in.ResourceARN, toDriverTags(in.Tags)); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +// untagResource handles POST /untagResource, removing tag keys from a resource. +func (h *Handler) untagResource(w http.ResponseWriter, r *http.Request) { + var in untagResourceRequest + if !decodeJSON(w, r, &in) { + return + } + + if err := h.bedrock.UntagResource(r.Context(), in.ResourceARN, in.TagKeys); err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, struct{}{}) +} + +// listTagsForResource handles POST /listTagsForResource, returning a resource's +// tags. +func (h *Handler) listTagsForResource(w http.ResponseWriter, r *http.Request) { + var in listTagsRequest + if !decodeJSON(w, r, &in) { + return + } + + tags, err := h.bedrock.ListTagsForResource(r.Context(), in.ResourceARN) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, listTagsResponse{Tags: toWireTags(tags)}) +} + +// toDriverTags converts wire tag pairs to driver tags. +func toDriverTags(pairs []tagPair) []bedrockdriver.Tag { + out := make([]bedrockdriver.Tag, 0, len(pairs)) + for _, p := range pairs { + out = append(out, bedrockdriver.Tag{Key: p.Key, Value: p.Value}) + } + + return out +} + +// toWireTags converts driver tags to wire tag pairs (always non-nil). +func toWireTags(tags []bedrockdriver.Tag) []tagPair { + out := make([]tagPair, 0, len(tags)) + for _, t := range tags { + out = append(out, tagPair{Key: t.Key, Value: t.Value}) + } + + return out +} diff --git a/server/aws/bedrock/types.go b/server/aws/bedrock/types.go index fc8677e1..041a7ad6 100644 --- a/server/aws/bedrock/types.go +++ b/server/aws/bedrock/types.go @@ -170,3 +170,89 @@ type converseResponse struct { Usage converseUsage `json:"usage"` Metrics converseMetrics `json:"metrics"` } + +// Tagging wire shapes. tagPair (json "key"/"value") is defined in management.go. + +type tagResourceRequest struct { + ResourceARN string `json:"resourceARN"` + Tags []tagPair `json:"tags"` +} + +type untagResourceRequest struct { + ResourceARN string `json:"resourceARN"` + TagKeys []string `json:"tagKeys"` +} + +type listTagsRequest struct { + ResourceARN string `json:"resourceARN"` +} + +type listTagsResponse struct { + Tags []tagPair `json:"tags"` +} + +// CountTokens wire shapes. The input is a union carrying exactly one of +// converse or invokeModel. + +type countTokensConverse struct { + Messages []converseMessage `json:"messages"` + System []converseTextBlock `json:"system"` +} + +type countTokensInvokeModel struct { + // Body is the model-native InvokeModel payload. The SDK serializes it as a + // base64 blob (a JSON string), so it is typed []byte — encoding/json + // base64-decodes it back to the raw bytes on unmarshal. + Body []byte `json:"body"` +} + +type countTokensRequest struct { + Converse *countTokensConverse `json:"converse"` + InvokeModel *countTokensInvokeModel `json:"invokeModel"` +} + +type countTokensResponse struct { + InputTokens int `json:"inputTokens"` +} + +// ApplyGuardrail wire shapes. guardrailIdentifier and guardrailVersion are path +// parameters; the body carries source, content, and outputScope. + +type applyGuardrailTextBlock struct { + Text string `json:"text"` + Qualifiers []string `json:"qualifiers,omitempty"` +} + +type applyGuardrailContentBlock struct { + Text *applyGuardrailTextBlock `json:"text,omitempty"` +} + +type applyGuardrailRequest struct { + Source string `json:"source"` + Content []applyGuardrailContentBlock `json:"content"` + OutputScope string `json:"outputScope,omitempty"` +} + +type applyGuardrailUsage struct { + TopicPolicyUnits int `json:"topicPolicyUnits"` + ContentPolicyUnits int `json:"contentPolicyUnits"` + WordPolicyUnits int `json:"wordPolicyUnits"` + SensitiveInformationPolicyUnits int `json:"sensitiveInformationPolicyUnits"` + SensitiveInformationPolicyFreeUnits int `json:"sensitiveInformationPolicyFreeUnits"` + ContextualGroundingPolicyUnits int `json:"contextualGroundingPolicyUnits"` +} + +type applyGuardrailOutputContent struct { + Text string `json:"text"` +} + +// applyGuardrailAssessment is emitted as an empty object list; the emulator +// performs no policy assessments. +type applyGuardrailAssessment struct{} + +type applyGuardrailResponse struct { + Usage applyGuardrailUsage `json:"usage"` + Action string `json:"action"` + Outputs []applyGuardrailOutputContent `json:"outputs"` + Assessments []applyGuardrailAssessment `json:"assessments"` +} diff --git a/server/aws/bedrockagent/agents.go b/server/aws/bedrockagent/agents.go new file mode 100644 index 00000000..1544d38b --- /dev/null +++ b/server/aws/bedrockagent/agents.go @@ -0,0 +1,225 @@ +package bedrockagent + +import ( + "net/http" + + badriver "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// serveAgents dispatches the /agents/ subtree. +func (h *Handler) serveAgents(w http.ResponseWriter, r *http.Request, segs []string) { + switch { + case len(segs) == 0: + h.serveAgentCollection(w, r) + case len(segs) == 1: + h.serveAgentItem(w, r, segs[0]) + case len(segs) == 2 && segs[1] == segAgentAliases: + h.serveAgentAlias(w, r, segs[0]) + default: + notFound(w, r.URL.Path) + } +} + +func (h *Handler) serveAgentCollection(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPut: + h.createAgent(w, r) + case http.MethodPost: + h.listAgents(w, r) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) serveAgentItem(w http.ResponseWriter, r *http.Request, id string) { + switch r.Method { + case http.MethodGet: + h.getAgent(w, r, id) + case http.MethodPut: + h.updateAgent(w, r, id) + case http.MethodDelete: + h.deleteAgent(w, r, id) + case http.MethodPost: + h.prepareAgent(w, r, id) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) serveAgentAlias(w http.ResponseWriter, r *http.Request, agentID string) { + if r.Method != http.MethodPut { + methodNotAllowed(w) + + return + } + + h.createAgentAlias(w, r, agentID) +} + +// --- operations --- + +func (h *Handler) createAgent(w http.ResponseWriter, r *http.Request) { + var in createAgentRequest + if !decodeJSON(w, r, &in) { + return + } + + agent, err := h.agent.CreateAgent(r.Context(), badriver.AgentConfig{ + Name: in.AgentName, + ResourceRoleArn: in.AgentResourceRoleArn, + FoundationModel: in.FoundationModel, + Instruction: in.Instruction, + Description: in.Description, + IdleSessionTTLInSeconds: in.IdleSessionTTLInSeconds, + ClientToken: in.ClientToken, + Tags: in.Tags, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, agentEnvelope{Agent: toAgentJSON(agent)}) +} + +func (h *Handler) getAgent(w http.ResponseWriter, r *http.Request, id string) { + agent, err := h.agent.GetAgent(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, agentEnvelope{Agent: toAgentJSON(agent)}) +} + +func (h *Handler) listAgents(w http.ResponseWriter, r *http.Request) { + agents, err := h.agent.ListAgents(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]agentSummaryJSON, 0, len(agents)) + for i := range agents { + out = append(out, toAgentSummaryJSON(&agents[i])) + } + + writeJSON(w, listAgentsResponse{AgentSummaries: out}) +} + +func (h *Handler) updateAgent(w http.ResponseWriter, r *http.Request, id string) { + var in createAgentRequest + if !decodeJSON(w, r, &in) { + return + } + + agent, err := h.agent.UpdateAgent(r.Context(), id, badriver.AgentConfig{ + Name: in.AgentName, + ResourceRoleArn: in.AgentResourceRoleArn, + FoundationModel: in.FoundationModel, + Instruction: in.Instruction, + Description: in.Description, + IdleSessionTTLInSeconds: in.IdleSessionTTLInSeconds, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, agentEnvelope{Agent: toAgentJSON(agent)}) +} + +func (h *Handler) deleteAgent(w http.ResponseWriter, r *http.Request, id string) { + status, err := h.agent.DeleteAgent(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, deleteAgentResponse{AgentID: id, AgentStatus: status}) +} + +func (h *Handler) prepareAgent(w http.ResponseWriter, r *http.Request, id string) { + agent, err := h.agent.PrepareAgent(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, prepareAgentResponse{ + AgentID: agent.ID, + AgentStatus: agent.Status, + AgentVersion: agent.Version, + PreparedAt: agent.PreparedAt, + }) +} + +func (h *Handler) createAgentAlias(w http.ResponseWriter, r *http.Request, agentID string) { + var in createAgentAliasRequest + if !decodeJSON(w, r, &in) { + return + } + + alias, err := h.agent.CreateAgentAlias(r.Context(), badriver.AgentAliasConfig{ + AgentID: agentID, + Name: in.AgentAliasName, + Description: in.Description, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, agentAliasEnvelope{AgentAlias: toAgentAliasJSON(alias)}) +} + +// --- converters --- + +func toAgentJSON(a *badriver.Agent) agentJSON { + return agentJSON{ + AgentID: a.ID, + AgentARN: a.ARN, + AgentName: a.Name, + AgentResourceRoleArn: a.ResourceRoleArn, + FoundationModel: a.FoundationModel, + Instruction: a.Instruction, + Description: a.Description, + AgentStatus: a.Status, + AgentVersion: a.Version, + IdleSessionTTLInSeconds: a.IdleSessionTTLInSeconds, + CreatedAt: a.CreatedAt, + UpdatedAt: a.UpdatedAt, + PreparedAt: a.PreparedAt, + } +} + +func toAgentSummaryJSON(a *badriver.Agent) agentSummaryJSON { + return agentSummaryJSON{ + AgentID: a.ID, + AgentName: a.Name, + AgentStatus: a.Status, + Description: a.Description, + UpdatedAt: a.UpdatedAt, + } +} + +func toAgentAliasJSON(a *badriver.AgentAlias) agentAliasJSON { + return agentAliasJSON{ + AgentAliasID: a.ID, + AgentAliasARN: a.ARN, + AgentAliasName: a.Name, + AgentID: a.AgentID, + AgentAliasStatus: a.Status, + Description: a.Description, + RoutingConfiguration: []string{}, + CreatedAt: a.CreatedAt, + UpdatedAt: a.UpdatedAt, + } +} diff --git a/server/aws/bedrockagent/errors.go b/server/aws/bedrockagent/errors.go new file mode 100644 index 00000000..b2133e07 --- /dev/null +++ b/server/aws/bedrockagent/errors.go @@ -0,0 +1,52 @@ +package bedrockagent + +import ( + "encoding/json" + "net/http" + + cerrors "github.com/stackshy/cloudemu/v2/errors" +) + +// errorBody is the JSON body shape Bedrock Agent returns for failures. The SDK +// reads the X-Amzn-ErrorType header to map to a typed exception and falls back +// to the body's __type field if absent. +type errorBody struct { + Type string `json:"__type"` + Message string `json:"message"` +} + +// writeError writes a restJson1 error response with the given HTTP status, +// Bedrock Agent-shaped error type, and message. +func writeError(w http.ResponseWriter, status int, errType, msg string) { + w.Header().Set("Content-Type", contentTypeJSON) + w.Header().Set("X-Amzn-Errortype", errType) + w.WriteHeader(status) + + _ = json.NewEncoder(w).Encode(errorBody{Type: errType, Message: msg}) +} + +// writeErr maps cloudemu canonical errors to Bedrock Agent-shaped responses. +func writeErr(w http.ResponseWriter, err error) { + switch { + case cerrors.IsNotFound(err): + writeError(w, http.StatusNotFound, "ResourceNotFoundException", err.Error()) + case cerrors.IsAlreadyExists(err): + writeError(w, http.StatusConflict, "ConflictException", err.Error()) + case cerrors.IsInvalidArgument(err): + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + case cerrors.IsFailedPrecondition(err): + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + case cerrors.IsThrottled(err): + writeError(w, http.StatusTooManyRequests, "ThrottlingException", err.Error()) + default: + writeError(w, http.StatusInternalServerError, "InternalServerException", err.Error()) + } +} + +func methodNotAllowed(w http.ResponseWriter) { + writeError(w, http.StatusMethodNotAllowed, "ValidationException", "method not allowed") +} + +func notFound(w http.ResponseWriter, p string) { + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported path: "+p) +} diff --git a/server/aws/bedrockagent/flows.go b/server/aws/bedrockagent/flows.go new file mode 100644 index 00000000..7ab44a74 --- /dev/null +++ b/server/aws/bedrockagent/flows.go @@ -0,0 +1,172 @@ +package bedrockagent + +import ( + "net/http" + + badriver "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// serveFlows dispatches the /flows subtree. +func (h *Handler) serveFlows(w http.ResponseWriter, r *http.Request, segs []string) { + switch { + case len(segs) == 0: + h.serveFlowCollection(w, r) + case len(segs) == 1: + h.serveFlowItem(w, r, segs[0]) + default: + notFound(w, r.URL.Path) + } +} + +func (h *Handler) serveFlowCollection(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + h.createFlow(w, r) + case http.MethodGet: + h.listFlows(w, r) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) serveFlowItem(w http.ResponseWriter, r *http.Request, id string) { + switch r.Method { + case http.MethodGet: + h.getFlow(w, r, id) + case http.MethodPut: + h.updateFlow(w, r, id) + case http.MethodDelete: + h.deleteFlow(w, r, id) + case http.MethodPost: + h.prepareFlow(w, r, id) + default: + methodNotAllowed(w) + } +} + +// --- operations --- + +func (h *Handler) createFlow(w http.ResponseWriter, r *http.Request) { + var in createFlowRequest + if !decodeJSON(w, r, &in) { + return + } + + flow, err := h.agent.CreateFlow(r.Context(), badriver.FlowConfig{ + Name: in.Name, + ExecutionRoleArn: in.ExecutionRoleArn, + Description: in.Description, + CustomerEncryptionKeyArn: in.CustomerEncryptionKeyArn, + Definition: in.Definition, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toFlowJSON(flow)) +} + +func (h *Handler) getFlow(w http.ResponseWriter, r *http.Request, id string) { + flow, err := h.agent.GetFlow(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toFlowJSON(flow)) +} + +func (h *Handler) listFlows(w http.ResponseWriter, r *http.Request) { + flows, err := h.agent.ListFlows(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]flowSummaryJSON, 0, len(flows)) + for i := range flows { + out = append(out, toFlowSummaryJSON(&flows[i])) + } + + writeJSON(w, listFlowsResponse{FlowSummaries: out}) +} + +//nolint:dupl // structurally similar to updatePrompt but operates on a distinct resource type. +func (h *Handler) updateFlow(w http.ResponseWriter, r *http.Request, id string) { + var in createFlowRequest + if !decodeJSON(w, r, &in) { + return + } + + flow, err := h.agent.UpdateFlow(r.Context(), id, badriver.FlowConfig{ + Name: in.Name, + ExecutionRoleArn: in.ExecutionRoleArn, + Description: in.Description, + CustomerEncryptionKeyArn: in.CustomerEncryptionKeyArn, + Definition: in.Definition, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toFlowJSON(flow)) +} + +func (h *Handler) deleteFlow(w http.ResponseWriter, r *http.Request, id string) { + fid, err := h.agent.DeleteFlow(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, deleteFlowResponse{ID: fid}) +} + +func (h *Handler) prepareFlow(w http.ResponseWriter, r *http.Request, id string) { + flow, err := h.agent.PrepareFlow(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, prepareFlowResponse{ID: flow.ID, Status: flow.Status}) +} + +// --- converters --- + +func toFlowJSON(f *badriver.Flow) flowJSON { + return flowJSON{ + Arn: f.ARN, + ID: f.ID, + Name: f.Name, + Status: f.Status, + Version: f.Version, + ExecutionRoleArn: f.ExecutionRoleArn, + Description: f.Description, + CustomerEncryptionKeyArn: f.CustomerEncryptionKeyArn, + Definition: f.Definition, + CreatedAt: f.CreatedAt, + UpdatedAt: f.UpdatedAt, + } +} + +func toFlowSummaryJSON(f *badriver.Flow) flowSummaryJSON { + return flowSummaryJSON{ + Arn: f.ARN, + ID: f.ID, + Name: f.Name, + Status: f.Status, + Version: f.Version, + Description: f.Description, + CreatedAt: f.CreatedAt, + UpdatedAt: f.UpdatedAt, + } +} diff --git a/server/aws/bedrockagent/handler.go b/server/aws/bedrockagent/handler.go new file mode 100644 index 00000000..2c769212 --- /dev/null +++ b/server/aws/bedrockagent/handler.go @@ -0,0 +1,154 @@ +// Package bedrockagent implements the AWS Bedrock Agent authoring restJson1 +// control-plane API as a server.Handler. Point the real +// aws-sdk-go-v2/service/bedrockagent client at a Server registered with this +// handler and the agent, knowledge-base, data-source, flow, and prompt +// authoring lifecycles work end-to-end against an in-memory driver. +// +// Routing is by (HTTP method, path-template); method disambiguates same-path +// operations. URL shapes follow what the SDK emits: +// +// PUT /agents/ — CreateAgent +// POST /agents/ — ListAgents +// GET /agents/{agentId}/ — GetAgent +// PUT /agents/{agentId}/ — UpdateAgent +// DELETE /agents/{agentId}/ — DeleteAgent +// POST /agents/{agentId}/ — PrepareAgent +// PUT /agents/{agentId}/agentaliases/ — CreateAgentAlias +// PUT /knowledgebases/ — CreateKnowledgeBase +// POST /knowledgebases/ — ListKnowledgeBases +// GET /knowledgebases/{id} — GetKnowledgeBase +// PUT /knowledgebases/{id} — UpdateKnowledgeBase +// DELETE /knowledgebases/{id} — DeleteKnowledgeBase +// PUT /knowledgebases/{kb}/datasources/ — CreateDataSource +// POST /knowledgebases/{kb}/datasources/ — ListDataSources +// GET /knowledgebases/{kb}/datasources/{ds} — GetDataSource +// PUT /knowledgebases/{kb}/datasources/{ds} — UpdateDataSource +// DELETE /knowledgebases/{kb}/datasources/{ds} — DeleteDataSource +// PUT /knowledgebases/{kb}/datasources/{ds}/ingestionjobs/ — StartIngestionJob +// POST /flows/ — CreateFlow +// GET /flows/ — ListFlows +// GET /flows/{id}/ — GetFlow +// PUT /flows/{id}/ — UpdateFlow +// DELETE /flows/{id}/ — DeleteFlow +// POST /flows/{id}/ — PrepareFlow +// POST /prompts/ — CreatePrompt +// GET /prompts/ — ListPrompts +// GET /prompts/{id}/ — GetPrompt +// PUT /prompts/{id}/ — UpdatePrompt +// DELETE /prompts/{id}/ — DeletePrompt +package bedrockagent + +import ( + "encoding/json" + "errors" + "io" + "net/http" + "strings" + + badriver "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +const ( + contentTypeJSON = "application/json" + maxBodyBytes = 5 << 20 + + prefixAgents = "/agents" + prefixKB = "/knowledgebases" + prefixFlows = "/flows" + prefixPrompts = "/prompts" + + segAgentAliases = "agentaliases" + segDataSources = "datasources" + segIngestionJobs = "ingestionjobs" + + // segment counts identifying the nested data-source path shapes. + dsCollectionSegments = 2 // {kb}/datasources + dsItemSegments = 3 // {kb}/datasources/{ds} + ingestionSegments = 4 // {kb}/datasources/{ds}/ingestionjobs +) + +// Handler serves AWS Bedrock Agent restJson1 requests against a driver. +type Handler struct { + agent badriver.BedrockAgent +} + +// New returns a Bedrock Agent handler backed by drv. +func New(drv badriver.BedrockAgent) *Handler { + return &Handler{agent: drv} +} + +// Matches claims the Bedrock Agent authoring URL prefixes. +func (*Handler) Matches(r *http.Request) bool { + p := r.URL.Path + + return strings.HasPrefix(p, prefixAgents+"/") || + strings.HasPrefix(p, prefixKB) || + strings.HasPrefix(p, prefixFlows) || + strings.HasPrefix(p, prefixPrompts) +} + +// ServeHTTP routes by URL prefix. +func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + p := r.URL.Path + + switch { + case strings.HasPrefix(p, prefixAgents+"/"): + h.serveAgents(w, r, segments(p, prefixAgents)) + case strings.HasPrefix(p, prefixKB): + h.serveKnowledgeBases(w, r, segments(p, prefixKB)) + case strings.HasPrefix(p, prefixFlows): + h.serveFlows(w, r, segments(p, prefixFlows)) + case strings.HasPrefix(p, prefixPrompts): + h.servePrompts(w, r, segments(p, prefixPrompts)) + default: + notFound(w, p) + } +} + +// segments trims prefix and surrounding slashes from path and splits the +// remainder into path segments. An empty remainder yields a nil slice. +func segments(path, prefix string) []string { + rest := strings.Trim(strings.TrimPrefix(path, prefix), "/") + if rest == "" { + return nil + } + + return strings.Split(rest, "/") +} + +func decodeJSON(w http.ResponseWriter, r *http.Request, v any) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) + + if err := json.NewDecoder(r.Body).Decode(v); err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", "invalid JSON: "+err.Error()) + + return false + } + + return true +} + +// decodeBody decodes the request body but tolerates an empty body (some ops +// such as PrepareAgent send no payload). +func decodeBody(w http.ResponseWriter, r *http.Request, v any) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) + + err := json.NewDecoder(r.Body).Decode(v) + if errors.Is(err, io.EOF) { + return true + } + + if err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", "invalid JSON: "+err.Error()) + + return false + } + + return true +} + +func writeJSON(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", contentTypeJSON) + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(v) +} diff --git a/server/aws/bedrockagent/knowledgebases.go b/server/aws/bedrockagent/knowledgebases.go new file mode 100644 index 00000000..92f44aba --- /dev/null +++ b/server/aws/bedrockagent/knowledgebases.go @@ -0,0 +1,346 @@ +package bedrockagent + +import ( + "net/http" + + badriver "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// serveKnowledgeBases dispatches the /knowledgebases subtree, including the +// nested data-source and ingestion-job paths. +func (h *Handler) serveKnowledgeBases(w http.ResponseWriter, r *http.Request, segs []string) { + switch { + case len(segs) == 0: + h.serveKBCollection(w, r) + case len(segs) == 1: + h.serveKBItem(w, r, segs[0]) + case segs[1] == segDataSources: + h.serveDataSources(w, r, segs) + default: + notFound(w, r.URL.Path) + } +} + +func (h *Handler) serveKBCollection(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPut: + h.createKnowledgeBase(w, r) + case http.MethodPost: + h.listKnowledgeBases(w, r) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) serveKBItem(w http.ResponseWriter, r *http.Request, id string) { + switch r.Method { + case http.MethodGet: + h.getKnowledgeBase(w, r, id) + case http.MethodPut: + h.updateKnowledgeBase(w, r, id) + case http.MethodDelete: + h.deleteKnowledgeBase(w, r, id) + default: + methodNotAllowed(w) + } +} + +// serveDataSources dispatches /knowledgebases/{kb}/datasources[/{ds}[/ingestionjobs/]]. +func (h *Handler) serveDataSources(w http.ResponseWriter, r *http.Request, segs []string) { + kbID := segs[0] + + switch { + case len(segs) == dsCollectionSegments: + h.serveDSCollection(w, r, kbID) + case len(segs) == dsItemSegments: + h.serveDSItem(w, r, kbID, segs[2]) + case len(segs) == ingestionSegments && segs[3] == segIngestionJobs: + h.serveIngestion(w, r, kbID, segs[2]) + default: + notFound(w, r.URL.Path) + } +} + +func (h *Handler) serveDSCollection(w http.ResponseWriter, r *http.Request, kbID string) { + switch r.Method { + case http.MethodPut: + h.createDataSource(w, r, kbID) + case http.MethodPost: + h.listDataSources(w, r, kbID) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) serveDSItem(w http.ResponseWriter, r *http.Request, kbID, dsID string) { + switch r.Method { + case http.MethodGet: + h.getDataSource(w, r, kbID, dsID) + case http.MethodPut: + h.updateDataSource(w, r, kbID, dsID) + case http.MethodDelete: + h.deleteDataSource(w, r, kbID, dsID) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) serveIngestion(w http.ResponseWriter, r *http.Request, kbID, dsID string) { + if r.Method != http.MethodPut { + methodNotAllowed(w) + + return + } + + h.startIngestionJob(w, r, kbID, dsID) +} + +// --- knowledge-base operations --- + +func (h *Handler) createKnowledgeBase(w http.ResponseWriter, r *http.Request) { + var in createKnowledgeBaseRequest + if !decodeJSON(w, r, &in) { + return + } + + kb, err := h.agent.CreateKnowledgeBase(r.Context(), badriver.KnowledgeBaseConfig{ + Name: in.Name, + RoleArn: in.RoleArn, + Description: in.Description, + KnowledgeBaseConfiguration: in.KnowledgeBaseConfiguration, + StorageConfiguration: in.StorageConfiguration, + Tags: in.Tags, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, knowledgeBaseEnvelope{KnowledgeBase: toKnowledgeBaseJSON(kb)}) +} + +func (h *Handler) getKnowledgeBase(w http.ResponseWriter, r *http.Request, id string) { + kb, err := h.agent.GetKnowledgeBase(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, knowledgeBaseEnvelope{KnowledgeBase: toKnowledgeBaseJSON(kb)}) +} + +func (h *Handler) listKnowledgeBases(w http.ResponseWriter, r *http.Request) { + kbs, err := h.agent.ListKnowledgeBases(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]knowledgeBaseSummaryJSON, 0, len(kbs)) + for i := range kbs { + out = append(out, toKnowledgeBaseSummaryJSON(&kbs[i])) + } + + writeJSON(w, listKnowledgeBasesResponse{KnowledgeBaseSummaries: out}) +} + +func (h *Handler) updateKnowledgeBase(w http.ResponseWriter, r *http.Request, id string) { + var in createKnowledgeBaseRequest + if !decodeJSON(w, r, &in) { + return + } + + kb, err := h.agent.UpdateKnowledgeBase(r.Context(), id, badriver.KnowledgeBaseConfig{ + Name: in.Name, + RoleArn: in.RoleArn, + Description: in.Description, + KnowledgeBaseConfiguration: in.KnowledgeBaseConfiguration, + StorageConfiguration: in.StorageConfiguration, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, knowledgeBaseEnvelope{KnowledgeBase: toKnowledgeBaseJSON(kb)}) +} + +func (h *Handler) deleteKnowledgeBase(w http.ResponseWriter, r *http.Request, id string) { + status, err := h.agent.DeleteKnowledgeBase(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, deleteKnowledgeBaseResponse{KnowledgeBaseID: id, Status: status}) +} + +// --- data-source operations --- + +func (h *Handler) createDataSource(w http.ResponseWriter, r *http.Request, kbID string) { + var in createDataSourceRequest + if !decodeJSON(w, r, &in) { + return + } + + ds, err := h.agent.CreateDataSource(r.Context(), badriver.DataSourceConfig{ + KnowledgeBaseID: kbID, + Name: in.Name, + Description: in.Description, + DataDeletionPolicy: in.DataDeletionPolicy, + DataSourceConfiguration: in.DataSourceConfiguration, + VectorIngestionConfiguration: in.VectorIngestionConfiguration, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, dataSourceEnvelope{DataSource: toDataSourceJSON(ds)}) +} + +func (h *Handler) getDataSource(w http.ResponseWriter, r *http.Request, kbID, dsID string) { + ds, err := h.agent.GetDataSource(r.Context(), kbID, dsID) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, dataSourceEnvelope{DataSource: toDataSourceJSON(ds)}) +} + +func (h *Handler) listDataSources(w http.ResponseWriter, r *http.Request, kbID string) { + dss, err := h.agent.ListDataSources(r.Context(), kbID) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]dataSourceSummaryJSON, 0, len(dss)) + for i := range dss { + out = append(out, toDataSourceSummaryJSON(&dss[i])) + } + + writeJSON(w, listDataSourcesResponse{DataSourceSummaries: out}) +} + +func (h *Handler) updateDataSource(w http.ResponseWriter, r *http.Request, kbID, dsID string) { + var in createDataSourceRequest + if !decodeJSON(w, r, &in) { + return + } + + ds, err := h.agent.UpdateDataSource(r.Context(), badriver.DataSourceConfig{ + KnowledgeBaseID: kbID, + Name: in.Name, + Description: in.Description, + DataDeletionPolicy: in.DataDeletionPolicy, + DataSourceConfiguration: in.DataSourceConfiguration, + }, dsID) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, dataSourceEnvelope{DataSource: toDataSourceJSON(ds)}) +} + +func (h *Handler) deleteDataSource(w http.ResponseWriter, r *http.Request, kbID, dsID string) { + status, err := h.agent.DeleteDataSource(r.Context(), kbID, dsID) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, deleteDataSourceResponse{DataSourceID: dsID, KnowledgeBaseID: kbID, Status: status}) +} + +func (h *Handler) startIngestionJob(w http.ResponseWriter, r *http.Request, kbID, dsID string) { + var in startIngestionJobRequest + if !decodeBody(w, r, &in) { + return + } + + job, err := h.agent.StartIngestionJob(r.Context(), kbID, dsID, in.Description) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, ingestionJobEnvelope{IngestionJob: toIngestionJobJSON(job)}) +} + +// --- converters --- + +func toKnowledgeBaseJSON(kb *badriver.KnowledgeBase) knowledgeBaseJSON { + return knowledgeBaseJSON{ + KnowledgeBaseID: kb.ID, + KnowledgeBaseARN: kb.ARN, + Name: kb.Name, + RoleArn: kb.RoleArn, + Description: kb.Description, + Status: kb.Status, + KnowledgeBaseConfiguration: kb.KnowledgeBaseConfiguration, + StorageConfiguration: kb.StorageConfiguration, + CreatedAt: kb.CreatedAt, + UpdatedAt: kb.UpdatedAt, + } +} + +func toKnowledgeBaseSummaryJSON(kb *badriver.KnowledgeBase) knowledgeBaseSummaryJSON { + return knowledgeBaseSummaryJSON{ + KnowledgeBaseID: kb.ID, + Name: kb.Name, + Status: kb.Status, + Description: kb.Description, + UpdatedAt: kb.UpdatedAt, + } +} + +func toDataSourceJSON(ds *badriver.DataSource) dataSourceJSON { + return dataSourceJSON{ + DataSourceID: ds.ID, + KnowledgeBaseID: ds.KnowledgeBaseID, + Name: ds.Name, + Description: ds.Description, + Status: ds.Status, + DataDeletionPolicy: ds.DataDeletionPolicy, + DataSourceConfiguration: ds.DataSourceConfiguration, + CreatedAt: ds.CreatedAt, + UpdatedAt: ds.UpdatedAt, + } +} + +func toDataSourceSummaryJSON(ds *badriver.DataSource) dataSourceSummaryJSON { + return dataSourceSummaryJSON{ + DataSourceID: ds.ID, + KnowledgeBaseID: ds.KnowledgeBaseID, + Name: ds.Name, + Status: ds.Status, + Description: ds.Description, + UpdatedAt: ds.UpdatedAt, + } +} + +func toIngestionJobJSON(j *badriver.IngestionJob) ingestionJobJSON { + return ingestionJobJSON{ + IngestionJobID: j.ID, + KnowledgeBaseID: j.KnowledgeBaseID, + DataSourceID: j.DataSourceID, + Description: j.Description, + Status: j.Status, + StartedAt: j.StartedAt, + UpdatedAt: j.UpdatedAt, + } +} diff --git a/server/aws/bedrockagent/prompts.go b/server/aws/bedrockagent/prompts.go new file mode 100644 index 00000000..eb6e5649 --- /dev/null +++ b/server/aws/bedrockagent/prompts.go @@ -0,0 +1,157 @@ +package bedrockagent + +import ( + "net/http" + + badriver "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +// servePrompts dispatches the /prompts subtree. +func (h *Handler) servePrompts(w http.ResponseWriter, r *http.Request, segs []string) { + switch { + case len(segs) == 0: + h.servePromptCollection(w, r) + case len(segs) == 1: + h.servePromptItem(w, r, segs[0]) + default: + notFound(w, r.URL.Path) + } +} + +func (h *Handler) servePromptCollection(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + h.createPrompt(w, r) + case http.MethodGet: + h.listPrompts(w, r) + default: + methodNotAllowed(w) + } +} + +func (h *Handler) servePromptItem(w http.ResponseWriter, r *http.Request, id string) { + switch r.Method { + case http.MethodGet: + h.getPrompt(w, r, id) + case http.MethodPut: + h.updatePrompt(w, r, id) + case http.MethodDelete: + h.deletePrompt(w, r, id) + default: + methodNotAllowed(w) + } +} + +// --- operations --- + +func (h *Handler) createPrompt(w http.ResponseWriter, r *http.Request) { + var in createPromptRequest + if !decodeJSON(w, r, &in) { + return + } + + prompt, err := h.agent.CreatePrompt(r.Context(), badriver.PromptConfig{ + Name: in.Name, + Description: in.Description, + DefaultVariant: in.DefaultVariant, + CustomerEncryptionKeyArn: in.CustomerEncryptionKeyArn, + Variants: in.Variants, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toPromptJSON(prompt)) +} + +func (h *Handler) getPrompt(w http.ResponseWriter, r *http.Request, id string) { + prompt, err := h.agent.GetPrompt(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toPromptJSON(prompt)) +} + +func (h *Handler) listPrompts(w http.ResponseWriter, r *http.Request) { + prompts, err := h.agent.ListPrompts(r.Context()) + if err != nil { + writeErr(w, err) + + return + } + + out := make([]promptSummaryJSON, 0, len(prompts)) + for i := range prompts { + out = append(out, toPromptSummaryJSON(&prompts[i])) + } + + writeJSON(w, listPromptsResponse{PromptSummaries: out}) +} + +//nolint:dupl // structurally similar to updateFlow but operates on a distinct resource type. +func (h *Handler) updatePrompt(w http.ResponseWriter, r *http.Request, id string) { + var in createPromptRequest + if !decodeJSON(w, r, &in) { + return + } + + prompt, err := h.agent.UpdatePrompt(r.Context(), id, badriver.PromptConfig{ + Name: in.Name, + Description: in.Description, + DefaultVariant: in.DefaultVariant, + CustomerEncryptionKeyArn: in.CustomerEncryptionKeyArn, + Variants: in.Variants, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toPromptJSON(prompt)) +} + +func (h *Handler) deletePrompt(w http.ResponseWriter, r *http.Request, id string) { + pid, err := h.agent.DeletePrompt(r.Context(), id) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, deletePromptResponse{ID: pid}) +} + +// --- converters --- + +func toPromptJSON(p *badriver.Prompt) promptJSON { + return promptJSON{ + Arn: p.ARN, + ID: p.ID, + Name: p.Name, + Version: p.Version, + Description: p.Description, + DefaultVariant: p.DefaultVariant, + CustomerEncryptionKeyArn: p.CustomerEncryptionKeyArn, + Variants: p.Variants, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + } +} + +func toPromptSummaryJSON(p *badriver.Prompt) promptSummaryJSON { + return promptSummaryJSON{ + Arn: p.ARN, + ID: p.ID, + Name: p.Name, + Version: p.Version, + Description: p.Description, + CreatedAt: p.CreatedAt, + UpdatedAt: p.UpdatedAt, + } +} diff --git a/server/aws/bedrockagent/sdk_roundtrip_test.go b/server/aws/bedrockagent/sdk_roundtrip_test.go new file mode 100644 index 00000000..e9a59d9b --- /dev/null +++ b/server/aws/bedrockagent/sdk_roundtrip_test.go @@ -0,0 +1,358 @@ +package bedrockagent_test + +import ( + "context" + "errors" + "net/http/httptest" + "testing" + + "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + awsba "github.com/aws/aws-sdk-go-v2/service/bedrockagent" + batypes "github.com/aws/aws-sdk-go-v2/service/bedrockagent/types" + + "github.com/stackshy/cloudemu/v2/config" + providerba "github.com/stackshy/cloudemu/v2/providers/aws/bedrockagent" + serverba "github.com/stackshy/cloudemu/v2/server/aws/bedrockagent" +) + +const ( + roleArn = "arn:aws:iam::123456789012:role/bedrock" + embedArn = "arn:aws:bedrock:us-east-1::foundation-model/amazon.titan-embed-text-v1" + bucketArn = "arn:aws:s3:::my-kb-bucket" + claudeModel = "anthropic.claude-3-sonnet-20240229-v1:0" +) + +func newClient(t *testing.T) *awsba.Client { + t.Helper() + + drv := providerba.New(config.NewOptions()) + ts := httptest.NewServer(serverba.New(drv)) + t.Cleanup(ts.Close) + + cfg, err := awsconfig.LoadDefaultConfig(context.Background(), + awsconfig.WithRegion("us-east-1"), + awsconfig.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("test", "test", "")), + ) + if err != nil { + t.Fatalf("aws config: %v", err) + } + + return awsba.NewFromConfig(cfg, func(o *awsba.Options) { + o.BaseEndpoint = aws.String(ts.URL) + }) +} + +func TestSDKAgentLifecycle(t *testing.T) { + client := newClient(t) + ctx := context.Background() + + created, err := client.CreateAgent(ctx, &awsba.CreateAgentInput{ + AgentName: aws.String("my-agent"), + AgentResourceRoleArn: aws.String(roleArn), + FoundationModel: aws.String(claudeModel), + Instruction: aws.String("You are a helpful assistant that answers questions."), + }) + if err != nil { + t.Fatalf("CreateAgent: %v", err) + } + + agentID := aws.ToString(created.Agent.AgentId) + if agentID == "" || aws.ToString(created.Agent.AgentArn) == "" { + t.Fatalf("expected agent id + arn, got %+v", created.Agent) + } + + if created.Agent.AgentStatus != batypes.AgentStatusNotPrepared { + t.Fatalf("got status %q, want NOT_PREPARED", created.Agent.AgentStatus) + } + + got, err := client.GetAgent(ctx, &awsba.GetAgentInput{AgentId: aws.String(agentID)}) + if err != nil { + t.Fatalf("GetAgent: %v", err) + } + + if aws.ToString(got.Agent.AgentName) != "my-agent" { + t.Fatalf("got name %q", aws.ToString(got.Agent.AgentName)) + } + + prep, err := client.PrepareAgent(ctx, &awsba.PrepareAgentInput{AgentId: aws.String(agentID)}) + if err != nil { + t.Fatalf("PrepareAgent: %v", err) + } + + if prep.AgentStatus != batypes.AgentStatusPrepared || prep.PreparedAt == nil { + t.Fatalf("expected PREPARED with preparedAt, got %+v", prep) + } + + list, err := client.ListAgents(ctx, &awsba.ListAgentsInput{}) + if err != nil { + t.Fatalf("ListAgents: %v", err) + } + + if len(list.AgentSummaries) != 1 { + t.Fatalf("got %d agent summaries, want 1", len(list.AgentSummaries)) + } + + upd, err := client.UpdateAgent(ctx, &awsba.UpdateAgentInput{ + AgentId: aws.String(agentID), + AgentName: aws.String("my-agent-renamed"), + AgentResourceRoleArn: aws.String(roleArn), + FoundationModel: aws.String(claudeModel), + }) + if err != nil { + t.Fatalf("UpdateAgent: %v", err) + } + + if aws.ToString(upd.Agent.AgentName) != "my-agent-renamed" { + t.Fatalf("update did not rename agent: %q", aws.ToString(upd.Agent.AgentName)) + } + + alias, err := client.CreateAgentAlias(ctx, &awsba.CreateAgentAliasInput{ + AgentId: aws.String(agentID), + AgentAliasName: aws.String("prod"), + }) + if err != nil { + t.Fatalf("CreateAgentAlias: %v", err) + } + + if aws.ToString(alias.AgentAlias.AgentAliasId) == "" { + t.Fatalf("expected alias id, got %+v", alias.AgentAlias) + } + + if _, err = client.DeleteAgent(ctx, &awsba.DeleteAgentInput{AgentId: aws.String(agentID)}); err != nil { + t.Fatalf("DeleteAgent: %v", err) + } + + _, err = client.GetAgent(ctx, &awsba.GetAgentInput{AgentId: aws.String(agentID)}) + assertNotFound(t, err) +} + +func TestSDKAgentNotFound(t *testing.T) { + client := newClient(t) + + _, err := client.GetAgent(context.Background(), &awsba.GetAgentInput{AgentId: aws.String("MISSING123")}) + assertNotFound(t, err) +} + +func TestSDKKnowledgeBaseAndDataSourceLifecycle(t *testing.T) { + client := newClient(t) + ctx := context.Background() + + kbCfg := &batypes.KnowledgeBaseConfiguration{ + Type: batypes.KnowledgeBaseTypeVector, + VectorKnowledgeBaseConfiguration: &batypes.VectorKnowledgeBaseConfiguration{ + EmbeddingModelArn: aws.String(embedArn), + }, + } + + kb, err := client.CreateKnowledgeBase(ctx, &awsba.CreateKnowledgeBaseInput{ + Name: aws.String("my-kb"), + RoleArn: aws.String(roleArn), + KnowledgeBaseConfiguration: kbCfg, + }) + if err != nil { + t.Fatalf("CreateKnowledgeBase: %v", err) + } + + kbID := aws.ToString(kb.KnowledgeBase.KnowledgeBaseId) + if kb.KnowledgeBase.Status != batypes.KnowledgeBaseStatusActive { + t.Fatalf("got status %q, want ACTIVE", kb.KnowledgeBase.Status) + } + + if kb.KnowledgeBase.KnowledgeBaseConfiguration == nil || + kb.KnowledgeBase.KnowledgeBaseConfiguration.Type != batypes.KnowledgeBaseTypeVector { + t.Fatalf("knowledge base configuration did not round-trip: %+v", kb.KnowledgeBase.KnowledgeBaseConfiguration) + } + + if _, err = client.GetKnowledgeBase(ctx, &awsba.GetKnowledgeBaseInput{ + KnowledgeBaseId: aws.String(kbID), + }); err != nil { + t.Fatalf("GetKnowledgeBase: %v", err) + } + + kbList, err := client.ListKnowledgeBases(ctx, &awsba.ListKnowledgeBasesInput{}) + if err != nil { + t.Fatalf("ListKnowledgeBases: %v", err) + } + + if len(kbList.KnowledgeBaseSummaries) != 1 { + t.Fatalf("got %d kb summaries, want 1", len(kbList.KnowledgeBaseSummaries)) + } + + testDataSourceLifecycle(t, client, kbID) + + if _, err = client.DeleteKnowledgeBase(ctx, &awsba.DeleteKnowledgeBaseInput{ + KnowledgeBaseId: aws.String(kbID), + }); err != nil { + t.Fatalf("DeleteKnowledgeBase: %v", err) + } + + _, err = client.GetKnowledgeBase(ctx, &awsba.GetKnowledgeBaseInput{KnowledgeBaseId: aws.String(kbID)}) + assertNotFound(t, err) +} + +func testDataSourceLifecycle(t *testing.T, client *awsba.Client, kbID string) { + t.Helper() + + ctx := context.Background() + dsCfg := &batypes.DataSourceConfiguration{ + Type: batypes.DataSourceTypeS3, + S3Configuration: &batypes.S3DataSourceConfiguration{BucketArn: aws.String(bucketArn)}, + } + + ds, err := client.CreateDataSource(ctx, &awsba.CreateDataSourceInput{ + KnowledgeBaseId: aws.String(kbID), + Name: aws.String("my-ds"), + DataSourceConfiguration: dsCfg, + }) + if err != nil { + t.Fatalf("CreateDataSource: %v", err) + } + + dsID := aws.ToString(ds.DataSource.DataSourceId) + if ds.DataSource.Status != batypes.DataSourceStatusAvailable { + t.Fatalf("got ds status %q, want AVAILABLE", ds.DataSource.Status) + } + + got, err := client.GetDataSource(ctx, &awsba.GetDataSourceInput{ + KnowledgeBaseId: aws.String(kbID), + DataSourceId: aws.String(dsID), + }) + if err != nil { + t.Fatalf("GetDataSource: %v", err) + } + + if got.DataSource.DataSourceConfiguration == nil || + got.DataSource.DataSourceConfiguration.Type != batypes.DataSourceTypeS3 { + t.Fatalf("data source configuration did not round-trip: %+v", got.DataSource.DataSourceConfiguration) + } + + job, err := client.StartIngestionJob(ctx, &awsba.StartIngestionJobInput{ + KnowledgeBaseId: aws.String(kbID), + DataSourceId: aws.String(dsID), + }) + if err != nil { + t.Fatalf("StartIngestionJob: %v", err) + } + + if job.IngestionJob.Status != batypes.IngestionJobStatusComplete { + t.Fatalf("got ingestion status %q, want COMPLETE", job.IngestionJob.Status) + } + + if _, err = client.DeleteDataSource(ctx, &awsba.DeleteDataSourceInput{ + KnowledgeBaseId: aws.String(kbID), + DataSourceId: aws.String(dsID), + }); err != nil { + t.Fatalf("DeleteDataSource: %v", err) + } +} + +func TestSDKFlowLifecycle(t *testing.T) { + client := newClient(t) + ctx := context.Background() + + created, err := client.CreateFlow(ctx, &awsba.CreateFlowInput{ + Name: aws.String("my-flow"), + ExecutionRoleArn: aws.String(roleArn), + Description: aws.String("a test flow"), + }) + if err != nil { + t.Fatalf("CreateFlow: %v", err) + } + + flowID := aws.ToString(created.Id) + if aws.ToString(created.Arn) == "" || created.Status != batypes.FlowStatusNotPrepared { + t.Fatalf("expected arn + NotPrepared, got %+v", created) + } + + got, err := client.GetFlow(ctx, &awsba.GetFlowInput{FlowIdentifier: aws.String(flowID)}) + if err != nil { + t.Fatalf("GetFlow: %v", err) + } + + if aws.ToString(got.Name) != "my-flow" { + t.Fatalf("got flow name %q", aws.ToString(got.Name)) + } + + prep, err := client.PrepareFlow(ctx, &awsba.PrepareFlowInput{FlowIdentifier: aws.String(flowID)}) + if err != nil { + t.Fatalf("PrepareFlow: %v", err) + } + + if prep.Status != batypes.FlowStatusPrepared { + t.Fatalf("got flow status %q, want Prepared", prep.Status) + } + + list, err := client.ListFlows(ctx, &awsba.ListFlowsInput{}) + if err != nil { + t.Fatalf("ListFlows: %v", err) + } + + if len(list.FlowSummaries) != 1 { + t.Fatalf("got %d flow summaries, want 1", len(list.FlowSummaries)) + } + + if _, err = client.DeleteFlow(ctx, &awsba.DeleteFlowInput{FlowIdentifier: aws.String(flowID)}); err != nil { + t.Fatalf("DeleteFlow: %v", err) + } + + _, err = client.GetFlow(ctx, &awsba.GetFlowInput{FlowIdentifier: aws.String(flowID)}) + assertNotFound(t, err) +} + +func TestSDKPromptLifecycle(t *testing.T) { + client := newClient(t) + ctx := context.Background() + + created, err := client.CreatePrompt(ctx, &awsba.CreatePromptInput{ + Name: aws.String("my-prompt"), + Description: aws.String("a test prompt"), + }) + if err != nil { + t.Fatalf("CreatePrompt: %v", err) + } + + promptID := aws.ToString(created.Id) + if aws.ToString(created.Arn) == "" || aws.ToString(created.Version) != "DRAFT" { + t.Fatalf("expected arn + DRAFT version, got %+v", created) + } + + got, err := client.GetPrompt(ctx, &awsba.GetPromptInput{PromptIdentifier: aws.String(promptID)}) + if err != nil { + t.Fatalf("GetPrompt: %v", err) + } + + if aws.ToString(got.Name) != "my-prompt" { + t.Fatalf("got prompt name %q", aws.ToString(got.Name)) + } + + list, err := client.ListPrompts(ctx, &awsba.ListPromptsInput{}) + if err != nil { + t.Fatalf("ListPrompts: %v", err) + } + + if len(list.PromptSummaries) != 1 { + t.Fatalf("got %d prompt summaries, want 1", len(list.PromptSummaries)) + } + + if _, err = client.DeletePrompt(ctx, &awsba.DeletePromptInput{PromptIdentifier: aws.String(promptID)}); err != nil { + t.Fatalf("DeletePrompt: %v", err) + } + + _, err = client.GetPrompt(ctx, &awsba.GetPromptInput{PromptIdentifier: aws.String(promptID)}) + assertNotFound(t, err) +} + +func assertNotFound(t *testing.T, err error) { + t.Helper() + + if err == nil { + t.Fatal("expected ResourceNotFoundException, got nil") + } + + var nfe *batypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + } +} diff --git a/server/aws/bedrockagent/types.go b/server/aws/bedrockagent/types.go new file mode 100644 index 00000000..1114e733 --- /dev/null +++ b/server/aws/bedrockagent/types.go @@ -0,0 +1,289 @@ +package bedrockagent + +import "encoding/json" + +// JSON wire shapes for the AWS Bedrock Agent restJson1 protocol. Field names +// use the exact camelCase keys the real aws-sdk-go-v2/service/bedrockagent +// client emits and expects. Nested configuration blocks are passed through +// verbatim as json.RawMessage. + +// --- agents --- + +type createAgentRequest struct { + AgentName string `json:"agentName"` + AgentResourceRoleArn string `json:"agentResourceRoleArn"` + FoundationModel string `json:"foundationModel"` + Instruction string `json:"instruction"` + Description string `json:"description"` + IdleSessionTTLInSeconds int32 `json:"idleSessionTTLInSeconds"` + ClientToken string `json:"clientToken"` + Tags map[string]string `json:"tags"` +} + +type agentJSON struct { + AgentID string `json:"agentId"` + AgentARN string `json:"agentArn"` + AgentName string `json:"agentName"` + AgentResourceRoleArn string `json:"agentResourceRoleArn,omitempty"` + FoundationModel string `json:"foundationModel,omitempty"` + Instruction string `json:"instruction,omitempty"` + Description string `json:"description,omitempty"` + AgentStatus string `json:"agentStatus"` + AgentVersion string `json:"agentVersion"` + IdleSessionTTLInSeconds int32 `json:"idleSessionTTLInSeconds"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` + PreparedAt string `json:"preparedAt,omitempty"` +} + +type agentSummaryJSON struct { + AgentID string `json:"agentId"` + AgentName string `json:"agentName"` + AgentStatus string `json:"agentStatus"` + Description string `json:"description,omitempty"` + UpdatedAt string `json:"updatedAt"` +} + +type agentEnvelope struct { + Agent agentJSON `json:"agent"` +} + +type listAgentsResponse struct { + AgentSummaries []agentSummaryJSON `json:"agentSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +type deleteAgentResponse struct { + AgentID string `json:"agentId"` + AgentStatus string `json:"agentStatus"` +} + +type prepareAgentResponse struct { + AgentID string `json:"agentId"` + AgentStatus string `json:"agentStatus"` + AgentVersion string `json:"agentVersion"` + PreparedAt string `json:"preparedAt"` +} + +type createAgentAliasRequest struct { + AgentAliasName string `json:"agentAliasName"` + Description string `json:"description"` +} + +type agentAliasJSON struct { + AgentAliasID string `json:"agentAliasId"` + AgentAliasARN string `json:"agentAliasArn"` + AgentAliasName string `json:"agentAliasName"` + AgentID string `json:"agentId"` + AgentAliasStatus string `json:"agentAliasStatus"` + Description string `json:"description,omitempty"` + RoutingConfiguration []string `json:"routingConfiguration"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type agentAliasEnvelope struct { + AgentAlias agentAliasJSON `json:"agentAlias"` +} + +// --- knowledge bases --- + +type createKnowledgeBaseRequest struct { + Name string `json:"name"` + RoleArn string `json:"roleArn"` + Description string `json:"description"` + KnowledgeBaseConfiguration json.RawMessage `json:"knowledgeBaseConfiguration"` + StorageConfiguration json.RawMessage `json:"storageConfiguration"` + Tags map[string]string `json:"tags"` +} + +type knowledgeBaseJSON struct { + KnowledgeBaseID string `json:"knowledgeBaseId"` + KnowledgeBaseARN string `json:"knowledgeBaseArn"` + Name string `json:"name"` + RoleArn string `json:"roleArn"` + Description string `json:"description,omitempty"` + Status string `json:"status"` + KnowledgeBaseConfiguration json.RawMessage `json:"knowledgeBaseConfiguration,omitempty"` + StorageConfiguration json.RawMessage `json:"storageConfiguration,omitempty"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type knowledgeBaseSummaryJSON struct { + KnowledgeBaseID string `json:"knowledgeBaseId"` + Name string `json:"name"` + Status string `json:"status"` + Description string `json:"description,omitempty"` + UpdatedAt string `json:"updatedAt"` +} + +type knowledgeBaseEnvelope struct { + KnowledgeBase knowledgeBaseJSON `json:"knowledgeBase"` +} + +type listKnowledgeBasesResponse struct { + KnowledgeBaseSummaries []knowledgeBaseSummaryJSON `json:"knowledgeBaseSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +type deleteKnowledgeBaseResponse struct { + KnowledgeBaseID string `json:"knowledgeBaseId"` + Status string `json:"status"` +} + +// --- data sources --- + +type createDataSourceRequest struct { + Name string `json:"name"` + Description string `json:"description"` + DataDeletionPolicy string `json:"dataDeletionPolicy"` + DataSourceConfiguration json.RawMessage `json:"dataSourceConfiguration"` + VectorIngestionConfiguration json.RawMessage `json:"vectorIngestionConfiguration"` +} + +type dataSourceJSON struct { + DataSourceID string `json:"dataSourceId"` + KnowledgeBaseID string `json:"knowledgeBaseId"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Status string `json:"status"` + DataDeletionPolicy string `json:"dataDeletionPolicy,omitempty"` + DataSourceConfiguration json.RawMessage `json:"dataSourceConfiguration,omitempty"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type dataSourceSummaryJSON struct { + DataSourceID string `json:"dataSourceId"` + KnowledgeBaseID string `json:"knowledgeBaseId"` + Name string `json:"name"` + Status string `json:"status"` + Description string `json:"description,omitempty"` + UpdatedAt string `json:"updatedAt"` +} + +type dataSourceEnvelope struct { + DataSource dataSourceJSON `json:"dataSource"` +} + +type listDataSourcesResponse struct { + DataSourceSummaries []dataSourceSummaryJSON `json:"dataSourceSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +type deleteDataSourceResponse struct { + DataSourceID string `json:"dataSourceId"` + KnowledgeBaseID string `json:"knowledgeBaseId"` + Status string `json:"status"` +} + +type startIngestionJobRequest struct { + Description string `json:"description"` +} + +type ingestionJobJSON struct { + IngestionJobID string `json:"ingestionJobId"` + KnowledgeBaseID string `json:"knowledgeBaseId"` + DataSourceID string `json:"dataSourceId"` + Description string `json:"description,omitempty"` + Status string `json:"status"` + StartedAt string `json:"startedAt"` + UpdatedAt string `json:"updatedAt"` +} + +type ingestionJobEnvelope struct { + IngestionJob ingestionJobJSON `json:"ingestionJob"` +} + +// --- flows (flat responses) --- + +type createFlowRequest struct { + Name string `json:"name"` + ExecutionRoleArn string `json:"executionRoleArn"` + Description string `json:"description"` + CustomerEncryptionKeyArn string `json:"customerEncryptionKeyArn"` + Definition json.RawMessage `json:"definition"` +} + +type flowJSON struct { + Arn string `json:"arn"` + ID string `json:"id"` + Name string `json:"name"` + Status string `json:"status"` + Version string `json:"version"` + ExecutionRoleArn string `json:"executionRoleArn"` + Description string `json:"description,omitempty"` + CustomerEncryptionKeyArn string `json:"customerEncryptionKeyArn,omitempty"` + Definition json.RawMessage `json:"definition,omitempty"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type flowSummaryJSON struct { + Arn string `json:"arn"` + ID string `json:"id"` + Name string `json:"name"` + Status string `json:"status"` + Version string `json:"version"` + Description string `json:"description,omitempty"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type listFlowsResponse struct { + FlowSummaries []flowSummaryJSON `json:"flowSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +type deleteFlowResponse struct { + ID string `json:"id"` +} + +type prepareFlowResponse struct { + ID string `json:"id"` + Status string `json:"status"` +} + +// --- prompts (flat responses) --- + +type createPromptRequest struct { + Name string `json:"name"` + Description string `json:"description"` + DefaultVariant string `json:"defaultVariant"` + CustomerEncryptionKeyArn string `json:"customerEncryptionKeyArn"` + Variants json.RawMessage `json:"variants"` +} + +type promptJSON struct { + Arn string `json:"arn"` + ID string `json:"id"` + Name string `json:"name"` + Version string `json:"version"` + Description string `json:"description,omitempty"` + DefaultVariant string `json:"defaultVariant,omitempty"` + CustomerEncryptionKeyArn string `json:"customerEncryptionKeyArn,omitempty"` + Variants json.RawMessage `json:"variants,omitempty"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type promptSummaryJSON struct { + Arn string `json:"arn"` + ID string `json:"id"` + Name string `json:"name"` + Version string `json:"version"` + Description string `json:"description,omitempty"` + CreatedAt string `json:"createdAt"` + UpdatedAt string `json:"updatedAt"` +} + +type listPromptsResponse struct { + PromptSummaries []promptSummaryJSON `json:"promptSummaries"` + NextToken string `json:"nextToken,omitempty"` +} + +type deletePromptResponse struct { + ID string `json:"id"` + Version string `json:"version,omitempty"` +} diff --git a/server/aws/bedrockagentruntime/errors.go b/server/aws/bedrockagentruntime/errors.go new file mode 100644 index 00000000..5a541fc8 --- /dev/null +++ b/server/aws/bedrockagentruntime/errors.go @@ -0,0 +1,48 @@ +package bedrockagentruntime + +import ( + "encoding/json" + "net/http" + + cerrors "github.com/stackshy/cloudemu/v2/errors" +) + +// errorBody is the JSON body shape the runtime returns for failures. The SDK +// reads the X-Amzn-ErrorType header to map to a typed exception and falls back +// to the body's type field if absent. +type errorBody struct { + Type string `json:"__type"` + Message string `json:"message"` +} + +// writeError writes a restJson1 error response with the given HTTP status, +// error type, and message. +func writeError(w http.ResponseWriter, status int, errType, msg string) { + w.Header().Set("Content-Type", contentTypeJSON) + w.Header().Set("X-Amzn-Errortype", errType) + w.WriteHeader(status) + + _ = json.NewEncoder(w).Encode(errorBody{Type: errType, Message: msg}) +} + +// writeErr maps cloudemu canonical errors to runtime-shaped error responses. +func writeErr(w http.ResponseWriter, err error) { + switch { + case cerrors.IsNotFound(err): + writeError(w, http.StatusNotFound, "ResourceNotFoundException", err.Error()) + case cerrors.IsAlreadyExists(err): + writeError(w, http.StatusConflict, "ConflictException", err.Error()) + case cerrors.IsInvalidArgument(err): + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + case cerrors.IsFailedPrecondition(err): + writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + case cerrors.IsThrottled(err): + writeError(w, http.StatusTooManyRequests, "ThrottlingException", err.Error()) + default: + writeError(w, http.StatusInternalServerError, "InternalServerException", err.Error()) + } +} + +func methodNotAllowed(w http.ResponseWriter) { + writeError(w, http.StatusMethodNotAllowed, "ValidationException", "method not allowed") +} diff --git a/server/aws/bedrockagentruntime/eventstream.go b/server/aws/bedrockagentruntime/eventstream.go new file mode 100644 index 00000000..e1373e53 --- /dev/null +++ b/server/aws/bedrockagentruntime/eventstream.go @@ -0,0 +1,55 @@ +package bedrockagentruntime + +import ( + "encoding/base64" + "net/http" + + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" + "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi" + + bedrockagentruntimedriver "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +// InvokeAgent output header names (see the SDK deserializer header bindings). +const ( + headerAgentSessionID = "x-amz-bedrock-agent-session-id" + headerAgentContentType = "x-amzn-bedrock-agent-content-type" +) + +// writeInvokeAgentStream writes the InvokeAgent response as an +// application/vnd.amazon.eventstream body: a single chunk event whose JSON +// payload {"bytes":""} decodes to the completion text. The SDK +// reconstructs the completion by concatenating each chunk's decoded bytes. +func writeInvokeAgentStream(w http.ResponseWriter, res *bedrockagentruntimedriver.InvokeAgentResult) { + ct := res.ContentType + if ct == "" { + ct = contentTypeJSON + } + + w.Header().Set("Content-Type", contentTypeEventStream) + w.Header().Set(headerAgentSessionID, res.SessionID) + w.Header().Set(headerAgentContentType, ct) + w.WriteHeader(http.StatusOK) + + enc := eventstream.NewEncoder() + flusher, _ := w.(http.Flusher) + + writeChunk(w, enc, []byte(res.Completion)) + + if flusher != nil { + flusher.Flush() + } +} + +// writeChunk encodes a single "chunk" event carrying textBytes. +func writeChunk(w http.ResponseWriter, enc *eventstream.Encoder, textBytes []byte) { + var h eventstream.Headers + + h.Set(eventstreamapi.MessageTypeHeader, eventstream.StringValue(eventstreamapi.EventMessageType)) + h.Set(eventstreamapi.EventTypeHeader, eventstream.StringValue("chunk")) + h.Set(eventstreamapi.ContentTypeHeader, eventstream.StringValue(contentTypeJSON)) + + payload := []byte(`{"bytes":"` + base64.StdEncoding.EncodeToString(textBytes) + `"}`) + + _ = enc.Encode(w, eventstream.Message{Headers: h, Payload: payload}) +} diff --git a/server/aws/bedrockagentruntime/handler.go b/server/aws/bedrockagentruntime/handler.go new file mode 100644 index 00000000..247305dd --- /dev/null +++ b/server/aws/bedrockagentruntime/handler.go @@ -0,0 +1,120 @@ +// Package bedrockagentruntime implements the AWS Bedrock Agent runtime +// (bedrock-agent-runtime) restJson1 data-plane API as a server.Handler. Point +// the real aws-sdk-go-v2/service/bedrockagentruntime client at a Server +// registered with this handler and InvokeAgent (an eventstream response), +// Retrieve, and RetrieveAndGenerate all work end-to-end against an in-memory +// driver. +// +// URL shapes follow what the SDK emits: +// +// POST /agents/{agentId}/agentAliases/{agentAliasId}/sessions/{sessionId}/text — InvokeAgent +// POST /knowledgebases/{knowledgeBaseId}/retrieve — Retrieve +// POST /retrieveAndGenerate — RetrieveAndGenerate +// +// The Matches predicate is intentionally SPECIFIC to these suffixes (.../text, +// .../retrieve) and the exact /retrieveAndGenerate path so it does not collide +// with a bedrock-agent CONTROL-PLANE handler that also owns /agents/ and +// /knowledgebases. The orchestrator registers this runtime handler before the +// control-plane handler so the specific matches win. +package bedrockagentruntime + +import ( + "net/http" + "strings" + + bedrockagentruntimedriver "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +const ( + contentTypeJSON = "application/json" + contentTypeEventStream = "application/vnd.amazon.eventstream" + maxBodyBytes = 5 << 20 + + prefixAgents = "/agents/" + prefixKnowledgeBases = "/knowledgebases/" + suffixText = "/text" + suffixRetrieve = "/retrieve" + pathRetrieveAndGen = "/retrieveAndGenerate" +) + +// Handler serves AWS Bedrock Agent runtime restJson1 requests against a driver. +type Handler struct { + runtime bedrockagentruntimedriver.BedrockAgentRuntime +} + +// New returns a Bedrock Agent runtime handler backed by drv. +func New(drv bedrockagentruntimedriver.BedrockAgentRuntime) *Handler { + return &Handler{runtime: drv} +} + +// isInvokeAgentPath reports whether p is an InvokeAgent path. +func isInvokeAgentPath(p string) bool { + return strings.HasPrefix(p, prefixAgents) && strings.HasSuffix(p, suffixText) +} + +// isRetrievePath reports whether p is a Retrieve path. +func isRetrievePath(p string) bool { + return strings.HasPrefix(p, prefixKnowledgeBases) && strings.HasSuffix(p, suffixRetrieve) +} + +// claims reports whether path p belongs to this handler. +func claims(p string) bool { + return isInvokeAgentPath(p) || isRetrievePath(p) || p == pathRetrieveAndGen +} + +// Matches claims only the Bedrock Agent runtime data-plane paths. Every runtime +// operation is POST, so requiring POST here prevents shadowing a control-plane +// GET/DELETE whose resource id happens to end in "text" or "retrieve". +func (*Handler) Matches(r *http.Request) bool { + return r.Method == http.MethodPost && claims(r.URL.Path) +} + +// ServeHTTP routes by URL shape. Every runtime path is POST-only. +func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { + p := r.URL.Path + + if r.Method != http.MethodPost { + methodNotAllowed(w) + + return + } + + switch { + case isInvokeAgentPath(p): + h.serveInvokeAgent(w, r, p) + case isRetrievePath(p): + h.serveRetrieve(w, r, p) + case p == pathRetrieveAndGen: + h.retrieveAndGenerate(w, r) + default: + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported path: "+p) + } +} + +// serveInvokeAgent parses the path +// /agents/{agentId}/agentAliases/{agentAliasId}/sessions/{sessionId}/text and +// dispatches to the InvokeAgent operation. +func (h *Handler) serveInvokeAgent(w http.ResponseWriter, r *http.Request, p string) { + parts := strings.Split(strings.Trim(p, "/"), "/") + if len(parts) != 7 || parts[0] != "agents" || parts[2] != "agentAliases" || + parts[4] != "sessions" || parts[6] != "text" { + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported InvokeAgent path") + + return + } + + h.invokeAgent(w, r, parts[1], parts[3], parts[5]) +} + +// serveRetrieve parses the path /knowledgebases/{knowledgeBaseId}/retrieve and +// dispatches to the Retrieve operation. +func (h *Handler) serveRetrieve(w http.ResponseWriter, r *http.Request, p string) { + parts := strings.Split(strings.Trim(p, "/"), "/") + if len(parts) != 3 || parts[0] != "knowledgebases" || parts[2] != "retrieve" { + writeError(w, http.StatusNotFound, "ResourceNotFoundException", "unsupported Retrieve path") + + return + } + + h.retrieve(w, r, parts[1]) +} diff --git a/server/aws/bedrockagentruntime/operations.go b/server/aws/bedrockagentruntime/operations.go new file mode 100644 index 00000000..c9bac610 --- /dev/null +++ b/server/aws/bedrockagentruntime/operations.go @@ -0,0 +1,133 @@ +package bedrockagentruntime + +import ( + "encoding/json" + "net/http" + + bedrockagentruntimedriver "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +func (h *Handler) invokeAgent(w http.ResponseWriter, r *http.Request, agentID, agentAliasID, sessionID string) { + var in invokeAgentRequest + if !decodeJSONAllowEmpty(w, r, &in) { + return + } + + res, err := h.runtime.InvokeAgent(r.Context(), bedrockagentruntimedriver.InvokeAgentInput{ + AgentID: agentID, + AgentAliasID: agentAliasID, + SessionID: sessionID, + InputText: in.InputText, + EnableTrace: in.EnableTrace, + EndSession: in.EndSession, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeInvokeAgentStream(w, res) +} + +func (h *Handler) retrieve(w http.ResponseWriter, r *http.Request, knowledgeBaseID string) { + var in retrieveRequest + if !decodeJSON(w, r, &in) { + return + } + + res, err := h.runtime.Retrieve(r.Context(), bedrockagentruntimedriver.RetrieveInput{ + KnowledgeBaseID: knowledgeBaseID, + QueryText: in.RetrievalQuery.Text, + NextToken: in.NextToken, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, toRetrieveResponse(res)) +} + +func (h *Handler) retrieveAndGenerate(w http.ResponseWriter, r *http.Request) { + var in retrieveAndGenerateRequest + if !decodeJSON(w, r, &in) { + return + } + + res, err := h.runtime.RetrieveAndGenerate(r.Context(), bedrockagentruntimedriver.RetrieveAndGenerateInput{ + InputText: in.Input.Text, + SessionID: in.SessionID, + }) + if err != nil { + writeErr(w, err) + + return + } + + writeJSON(w, retrieveAndGenerateResponse{ + Output: retrieveAndGenerateOutputBody{Text: res.Text}, + SessionID: res.SessionID, + Citations: []any{}, + }) +} + +// --- converters --- + +func toRetrieveResponse(res *bedrockagentruntimedriver.RetrieveResult) retrieveResponse { + out := make([]knowledgeBaseRetrievalResult, 0, len(res.Results)) + + for i := range res.Results { + r := &res.Results[i] + item := knowledgeBaseRetrievalResult{ + Content: retrievalResultContent{Type: "TEXT", Text: r.Text}, + Score: r.Score, + } + + if r.LocationURI != "" { + item.Location = &retrievalResultLocation{ + Type: "S3", + S3Location: &retrievalResultS3Location{URI: r.LocationURI}, + } + } + + out = append(out, item) + } + + return retrieveResponse{RetrievalResults: out, NextToken: res.NextToken} +} + +// --- helpers --- + +func decodeJSON(w http.ResponseWriter, r *http.Request, v any) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) + + if err := json.NewDecoder(r.Body).Decode(v); err != nil { + writeError(w, http.StatusBadRequest, "ValidationException", "invalid JSON: "+err.Error()) + + return false + } + + return true +} + +// decodeJSONAllowEmpty decodes a JSON body but tolerates an empty body, since +// InvokeAgent's body members are all optional. +func decodeJSONAllowEmpty(w http.ResponseWriter, r *http.Request, v any) bool { + r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) + + if err := json.NewDecoder(r.Body).Decode(v); err != nil && err.Error() != "EOF" { + writeError(w, http.StatusBadRequest, "ValidationException", "invalid JSON: "+err.Error()) + + return false + } + + return true +} + +func writeJSON(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", contentTypeJSON) + w.WriteHeader(http.StatusOK) + _ = json.NewEncoder(w).Encode(v) +} diff --git a/server/aws/bedrockagentruntime/sdk_roundtrip_test.go b/server/aws/bedrockagentruntime/sdk_roundtrip_test.go new file mode 100644 index 00000000..64750af6 --- /dev/null +++ b/server/aws/bedrockagentruntime/sdk_roundtrip_test.go @@ -0,0 +1,139 @@ +package bedrockagentruntime_test + +import ( + "context" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/aws/aws-sdk-go-v2/aws" + awsconfig "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/credentials" + agentruntime "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime" + agentruntimetypes "github.com/aws/aws-sdk-go-v2/service/bedrockagentruntime/types" + + "github.com/stackshy/cloudemu/v2/config" + providerbar "github.com/stackshy/cloudemu/v2/providers/aws/bedrockagentruntime" + serverbar "github.com/stackshy/cloudemu/v2/server/aws/bedrockagentruntime" + svcbar "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime" +) + +func newClient(t *testing.T) *agentruntime.Client { + t.Helper() + + fc := config.NewFakeClock(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)) + opts := config.NewOptions(config.WithClock(fc), config.WithRegion("us-east-1")) + + // Exercise the full stack: provider mock -> portable service -> handler. + svc := svcbar.NewBedrockAgentRuntime(providerbar.New(opts)) + ts := httptest.NewServer(serverbar.New(svc)) + t.Cleanup(ts.Close) + + cfg, err := awsconfig.LoadDefaultConfig(context.Background(), + awsconfig.WithRegion("us-east-1"), + awsconfig.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("test", "test", "")), + ) + if err != nil { + t.Fatalf("aws config: %v", err) + } + + return agentruntime.NewFromConfig(cfg, func(o *agentruntime.Options) { + o.BaseEndpoint = aws.String(ts.URL) + }) +} + +func TestSDKInvokeAgent(t *testing.T) { + client := newClient(t) + + out, err := client.InvokeAgent(context.Background(), &agentruntime.InvokeAgentInput{ + AgentId: aws.String("AGENT123"), + AgentAliasId: aws.String("ALIAS123"), + SessionId: aws.String("session-abc"), + InputText: aws.String("Tell me about Bedrock agents"), + }) + if err != nil { + t.Fatalf("InvokeAgent: %v", err) + } + + if aws.ToString(out.SessionId) != "session-abc" { + t.Fatalf("got session id %q, want session-abc", aws.ToString(out.SessionId)) + } + + stream := out.GetStream() + defer stream.Close() + + var completion strings.Builder + + for event := range stream.Events() { + if c, ok := event.(*agentruntimetypes.ResponseStreamMemberChunk); ok { + completion.Write(c.Value.Bytes) + } + } + + if err := stream.Err(); err != nil { + t.Fatalf("stream error: %v", err) + } + + got := completion.String() + if got == "" { + t.Fatal("expected a non-empty completion") + } + + if !strings.Contains(got, "Tell me about Bedrock agents") { + t.Fatalf("completion %q does not echo the prompt", got) + } +} + +func TestSDKRetrieve(t *testing.T) { + client := newClient(t) + + out, err := client.Retrieve(context.Background(), &agentruntime.RetrieveInput{ + KnowledgeBaseId: aws.String("KB123"), + RetrievalQuery: &agentruntimetypes.KnowledgeBaseQuery{ + Text: aws.String("what is a knowledge base"), + }, + }) + if err != nil { + t.Fatalf("Retrieve: %v", err) + } + + if len(out.RetrievalResults) == 0 { + t.Fatal("expected non-empty retrievalResults") + } + + for _, r := range out.RetrievalResults { + if r.Content == nil || aws.ToString(r.Content.Text) == "" { + t.Fatal("expected a result with content.text") + } + + if !strings.Contains(aws.ToString(r.Content.Text), "what is a knowledge base") { + t.Fatalf("result %q does not echo the query", aws.ToString(r.Content.Text)) + } + } +} + +func TestSDKRetrieveAndGenerate(t *testing.T) { + client := newClient(t) + + out, err := client.RetrieveAndGenerate(context.Background(), &agentruntime.RetrieveAndGenerateInput{ + Input: &agentruntimetypes.RetrieveAndGenerateInput{ + Text: aws.String("Summarize the docs"), + }, + }) + if err != nil { + t.Fatalf("RetrieveAndGenerate: %v", err) + } + + if out.Output == nil || aws.ToString(out.Output.Text) == "" { + t.Fatal("expected a non-empty output.text") + } + + if aws.ToString(out.SessionId) == "" { + t.Fatal("expected a session id") + } + + if !strings.Contains(aws.ToString(out.Output.Text), "Summarize the docs") { + t.Fatalf("output %q does not echo the input", aws.ToString(out.Output.Text)) + } +} diff --git a/server/aws/bedrockagentruntime/types.go b/server/aws/bedrockagentruntime/types.go new file mode 100644 index 00000000..fe023f51 --- /dev/null +++ b/server/aws/bedrockagentruntime/types.go @@ -0,0 +1,71 @@ +package bedrockagentruntime + +// JSON wire shapes for the AWS Bedrock Agent runtime restJson1 protocol. Field +// names use the exact camelCase keys the real aws-sdk-go-v2 +// bedrockagentruntime client emits and expects, so requests decode and +// responses deserialize unchanged. + +// --- InvokeAgent request (path params carry agent/alias/session ids) --- + +type invokeAgentRequest struct { + InputText string `json:"inputText"` + EnableTrace bool `json:"enableTrace"` + EndSession bool `json:"endSession"` +} + +// --- Retrieve --- + +type retrievalQuery struct { + Text string `json:"text"` +} + +type retrieveRequest struct { + RetrievalQuery retrievalQuery `json:"retrievalQuery"` + NextToken string `json:"nextToken"` +} + +type retrievalResultContent struct { + Type string `json:"type"` + Text string `json:"text"` +} + +type retrievalResultS3Location struct { + URI string `json:"uri"` +} + +type retrievalResultLocation struct { + Type string `json:"type"` + S3Location *retrievalResultS3Location `json:"s3Location,omitempty"` +} + +type knowledgeBaseRetrievalResult struct { + Content retrievalResultContent `json:"content"` + Location *retrievalResultLocation `json:"location,omitempty"` + Score float64 `json:"score"` +} + +type retrieveResponse struct { + RetrievalResults []knowledgeBaseRetrievalResult `json:"retrievalResults"` + NextToken string `json:"nextToken,omitempty"` +} + +// --- RetrieveAndGenerate --- + +type retrieveAndGenerateInputBody struct { + Text string `json:"text"` +} + +type retrieveAndGenerateRequest struct { + Input retrieveAndGenerateInputBody `json:"input"` + SessionID string `json:"sessionId"` +} + +type retrieveAndGenerateOutputBody struct { + Text string `json:"text"` +} + +type retrieveAndGenerateResponse struct { + Output retrieveAndGenerateOutputBody `json:"output"` + SessionID string `json:"sessionId"` + Citations []any `json:"citations"` +} diff --git a/services/bedrock/asyncinvoke_jobs.go b/services/bedrock/asyncinvoke_jobs.go new file mode 100644 index 00000000..853daafa --- /dev/null +++ b/services/bedrock/asyncinvoke_jobs.go @@ -0,0 +1,162 @@ +package bedrock + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- Async invocation --- + +// StartAsyncInvoke starts an asynchronous model invocation. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *Bedrock) StartAsyncInvoke(ctx context.Context, cfg driver.StartAsyncInvokeConfig) (*driver.AsyncInvoke, error) { + out, err := b.do(ctx, "StartAsyncInvoke", cfg.ModelID, func() (any, error) { return b.driver.StartAsyncInvoke(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.AsyncInvoke), nil +} + +// GetAsyncInvoke retrieves an async invocation by its invocation ARN. +func (b *Bedrock) GetAsyncInvoke(ctx context.Context, invocationARN string) (*driver.AsyncInvoke, error) { + out, err := b.do(ctx, "GetAsyncInvoke", invocationARN, func() (any, error) { + return b.driver.GetAsyncInvoke(ctx, invocationARN) + }) + if err != nil { + return nil, err + } + + return out.(*driver.AsyncInvoke), nil +} + +// ListAsyncInvokes lists all async invocations. +func (b *Bedrock) ListAsyncInvokes(ctx context.Context) ([]driver.AsyncInvoke, error) { + out, err := b.do(ctx, "ListAsyncInvokes", nil, func() (any, error) { return b.driver.ListAsyncInvokes(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.AsyncInvoke), nil +} + +// --- Model import jobs --- + +// CreateModelImportJob starts a custom-model import job. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *Bedrock) CreateModelImportJob(ctx context.Context, cfg driver.ModelImportJobConfig) (*driver.ModelImportJob, error) { + out, err := b.do(ctx, "CreateModelImportJob", cfg.JobName, func() (any, error) { + return b.driver.CreateModelImportJob(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.ModelImportJob), nil +} + +// GetModelImportJob retrieves an import job by name or ARN. +func (b *Bedrock) GetModelImportJob(ctx context.Context, jobIdentifier string) (*driver.ModelImportJob, error) { + out, err := b.do(ctx, "GetModelImportJob", jobIdentifier, func() (any, error) { + return b.driver.GetModelImportJob(ctx, jobIdentifier) + }) + if err != nil { + return nil, err + } + + return out.(*driver.ModelImportJob), nil +} + +// ListModelImportJobs lists all import jobs. +func (b *Bedrock) ListModelImportJobs(ctx context.Context) ([]driver.ModelImportJob, error) { + out, err := b.do(ctx, "ListModelImportJobs", nil, func() (any, error) { return b.driver.ListModelImportJobs(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.ModelImportJob), nil +} + +// --- Model copy jobs --- + +// CreateModelCopyJob starts a model-copy job. +func (b *Bedrock) CreateModelCopyJob(ctx context.Context, cfg driver.ModelCopyJobConfig) (*driver.ModelCopyJob, error) { + out, err := b.do(ctx, "CreateModelCopyJob", cfg.TargetModelName, func() (any, error) { + return b.driver.CreateModelCopyJob(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.ModelCopyJob), nil +} + +// GetModelCopyJob retrieves a copy job by its job ARN. +func (b *Bedrock) GetModelCopyJob(ctx context.Context, jobARN string) (*driver.ModelCopyJob, error) { + out, err := b.do(ctx, "GetModelCopyJob", jobARN, func() (any, error) { return b.driver.GetModelCopyJob(ctx, jobARN) }) + if err != nil { + return nil, err + } + + return out.(*driver.ModelCopyJob), nil +} + +// ListModelCopyJobs lists all copy jobs. +func (b *Bedrock) ListModelCopyJobs(ctx context.Context) ([]driver.ModelCopyJob, error) { + out, err := b.do(ctx, "ListModelCopyJobs", nil, func() (any, error) { return b.driver.ListModelCopyJobs(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.ModelCopyJob), nil +} + +// --- Evaluation jobs --- + +// CreateEvaluationJob starts a model-evaluation job. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *Bedrock) CreateEvaluationJob(ctx context.Context, cfg driver.EvaluationJobConfig) (*driver.EvaluationJob, error) { + out, err := b.do(ctx, "CreateEvaluationJob", cfg.JobName, func() (any, error) { + return b.driver.CreateEvaluationJob(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.EvaluationJob), nil +} + +// GetEvaluationJob retrieves an evaluation job by name or ARN. +func (b *Bedrock) GetEvaluationJob(ctx context.Context, jobIdentifier string) (*driver.EvaluationJob, error) { + out, err := b.do(ctx, "GetEvaluationJob", jobIdentifier, func() (any, error) { + return b.driver.GetEvaluationJob(ctx, jobIdentifier) + }) + if err != nil { + return nil, err + } + + return out.(*driver.EvaluationJob), nil +} + +// ListEvaluationJobs lists all evaluation jobs. +func (b *Bedrock) ListEvaluationJobs(ctx context.Context) ([]driver.EvaluationJob, error) { + out, err := b.do(ctx, "ListEvaluationJobs", nil, func() (any, error) { return b.driver.ListEvaluationJobs(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.EvaluationJob), nil +} + +// StopEvaluationJob transitions an evaluation job to the Stopped state. +func (b *Bedrock) StopEvaluationJob(ctx context.Context, jobIdentifier string) error { + _, err := b.do(ctx, "StopEvaluationJob", jobIdentifier, func() (any, error) { + return nil, b.driver.StopEvaluationJob(ctx, jobIdentifier) + }) + + return err +} diff --git a/services/bedrock/bedrock.go b/services/bedrock/bedrock.go index 348f2b30..1fe1e124 100644 --- a/services/bedrock/bedrock.go +++ b/services/bedrock/bedrock.go @@ -207,6 +207,58 @@ func (b *Bedrock) Converse(ctx context.Context, in driver.ConverseInput) (*drive return out.(*driver.ConverseOutput), nil } +// CountTokens returns an emulated input-token count for a would-be request. +// +//nolint:gocritic // in matches the driver interface signature; read without mutation. +func (b *Bedrock) CountTokens(ctx context.Context, in driver.CountTokensInput) (int, error) { + out, err := b.do(ctx, "CountTokens", in.ModelID, func() (any, error) { return b.driver.CountTokens(ctx, in) }) + if err != nil { + return 0, err + } + + return out.(int), nil +} + +// ApplyGuardrail evaluates content against a guardrail. +func (b *Bedrock) ApplyGuardrail(ctx context.Context, in driver.ApplyGuardrailInput) (*driver.ApplyGuardrailOutput, error) { + out, err := b.do(ctx, "ApplyGuardrail", in.GuardrailIdentifier, func() (any, error) { return b.driver.ApplyGuardrail(ctx, in) }) + if err != nil { + return nil, err + } + + return out.(*driver.ApplyGuardrailOutput), nil +} + +// TagResource attaches tags to a Bedrock resource by ARN. +func (b *Bedrock) TagResource(ctx context.Context, resourceARN string, tags []driver.Tag) error { + _, err := b.do(ctx, "TagResource", resourceARN, func() (any, error) { + return nil, b.driver.TagResource(ctx, resourceARN, tags) + }) + + return err +} + +// UntagResource removes tag keys from a Bedrock resource by ARN. +func (b *Bedrock) UntagResource(ctx context.Context, resourceARN string, tagKeys []string) error { + _, err := b.do(ctx, "UntagResource", resourceARN, func() (any, error) { + return nil, b.driver.UntagResource(ctx, resourceARN, tagKeys) + }) + + return err +} + +// ListTagsForResource lists the tags attached to a Bedrock resource by ARN. +func (b *Bedrock) ListTagsForResource(ctx context.Context, resourceARN string) ([]driver.Tag, error) { + out, err := b.do(ctx, "ListTagsForResource", resourceARN, func() (any, error) { + return b.driver.ListTagsForResource(ctx, resourceARN) + }) + if err != nil { + return nil, err + } + + return out.([]driver.Tag), nil +} + // CreateGuardrail creates a content guardrail. // //nolint:gocritic // cfg matches the driver interface signature; copied once on entry. @@ -229,9 +281,10 @@ func (b *Bedrock) GetGuardrail(ctx context.Context, identifier, version string) return out.(*driver.Guardrail), nil } -// ListGuardrails lists all guardrails. -func (b *Bedrock) ListGuardrails(ctx context.Context) ([]driver.Guardrail, error) { - out, err := b.do(ctx, "ListGuardrails", nil, func() (any, error) { return b.driver.ListGuardrails(ctx) }) +// ListGuardrails lists guardrails, optionally scoped to a single identifier +// (which expands to one entry per version). +func (b *Bedrock) ListGuardrails(ctx context.Context, identifier string) ([]driver.Guardrail, error) { + out, err := b.do(ctx, "ListGuardrails", identifier, func() (any, error) { return b.driver.ListGuardrails(ctx, identifier) }) if err != nil { return nil, err } @@ -251,9 +304,29 @@ func (b *Bedrock) UpdateGuardrail(ctx context.Context, identifier string, cfg dr return out.(*driver.Guardrail), nil } -// DeleteGuardrail deletes a guardrail. -func (b *Bedrock) DeleteGuardrail(ctx context.Context, identifier string) error { - _, err := b.do(ctx, "DeleteGuardrail", identifier, func() (any, error) { return nil, b.driver.DeleteGuardrail(ctx, identifier) }) +// CreateGuardrailVersion snapshots a guardrail's DRAFT into a new immutable +// numbered version, returning the guardrail ID and the assigned version. +func (b *Bedrock) CreateGuardrailVersion( + ctx context.Context, identifier, description string, +) (guardrailID, version string, err error) { + out, err := b.do(ctx, "CreateGuardrailVersion", identifier, func() (any, error) { + id, ver, cErr := b.driver.CreateGuardrailVersion(ctx, identifier, description) + return [2]string{id, ver}, cErr + }) + if err != nil { + return "", "", err + } + + pair := out.([2]string) + + return pair[0], pair[1], nil +} + +// DeleteGuardrail deletes a guardrail, or a single version when version is set. +func (b *Bedrock) DeleteGuardrail(ctx context.Context, identifier, version string) error { + _, err := b.do(ctx, "DeleteGuardrail", identifier, func() (any, error) { + return nil, b.driver.DeleteGuardrail(ctx, identifier, version) + }) return err } diff --git a/services/bedrock/bedrock_test.go b/services/bedrock/bedrock_test.go index ba0cf5fa..e89a67c7 100644 --- a/services/bedrock/bedrock_test.go +++ b/services/bedrock/bedrock_test.go @@ -166,11 +166,11 @@ func TestManagementWrappers(t *testing.T) { _, err = b.GetGuardrail(ctx, g.ID, "") require.NoError(t, err) - gs, err := b.ListGuardrails(ctx) + gs, err := b.ListGuardrails(ctx, "") require.NoError(t, err) require.Len(t, gs, 1) - require.NoError(t, b.DeleteGuardrail(ctx, g.ID)) + require.NoError(t, b.DeleteGuardrail(ctx, g.ID, "")) pt, err := b.CreateProvisionedModelThroughput(ctx, driver.ProvisionedThroughputConfig{ ProvisionedModelName: "pt-1", diff --git a/services/bedrock/driver/asyncinvoke_jobs.go b/services/bedrock/driver/asyncinvoke_jobs.go new file mode 100644 index 00000000..addb0048 --- /dev/null +++ b/services/bedrock/driver/asyncinvoke_jobs.go @@ -0,0 +1,129 @@ +package driver + +// Async invocation status values (bedrock-runtime). The set is deliberately +// smaller than the customization-job set: there is no "Submitted". +const ( + AsyncInProgress = "InProgress" + AsyncCompleted = "Completed" + AsyncFailed = "Failed" +) + +// Evaluation-job type values. +const ( + EvaluationTypeAutomated = "Automated" + EvaluationTypeHuman = "Human" +) + +// AsyncInvokeOutputConfig is the S3 delivery target for an async invocation's +// output, mirroring the s3OutputDataConfig union member. +type AsyncInvokeOutputConfig struct { + S3URI string + BucketOwner string + KMSKeyID string +} + +// StartAsyncInvokeConfig describes an asynchronous model invocation to start. +// ModelInput is a model-native JSON document carried through verbatim. +type StartAsyncInvokeConfig struct { + ClientRequestToken string + ModelID string + ModelInput []byte + Output AsyncInvokeOutputConfig + Tags map[string]string +} + +// AsyncInvoke describes an asynchronous model invocation. +type AsyncInvoke struct { + InvocationARN string + ModelARN string + ClientRequestToken string + Status string + Output AsyncInvokeOutputConfig + FailureMessage string + SubmitTime string + LastModifiedTime string + EndTime string +} + +// ModelImportJobConfig describes a custom-model import job to create. +type ModelImportJobConfig struct { + JobName string + ImportedModelName string + RoleARN string + ModelDataSourceS3URI string + ClientRequestToken string + ImportedModelKMSKeyID string + JobTags map[string]string + ImportedModelTags map[string]string +} + +// ModelImportJob describes a custom-model import job. +type ModelImportJob struct { + JobARN string + JobName string + ImportedModelName string + ImportedModelARN string + RoleARN string + ModelDataSourceS3URI string + Status string + FailureMessage string + CreationTime string + LastModifiedTime string + EndTime string +} + +// ModelCopyJobConfig describes a model-copy job to create. +type ModelCopyJobConfig struct { + SourceModelARN string + TargetModelName string + ClientRequestToken string + ModelKMSKeyID string + TargetModelTags map[string]string +} + +// ModelCopyJob describes a model-copy job. +type ModelCopyJob struct { + JobARN string + SourceAccountID string + SourceModelARN string + SourceModelName string + TargetModelName string + TargetModelARN string + TargetModelKMSKeyARN string + Status string + FailureMessage string + CreationTime string +} + +// EvaluationJobConfig describes an evaluation job to create. EvaluationConfig +// and InferenceConfig are opaque JSON documents carried through verbatim. +type EvaluationJobConfig struct { + JobName string + RoleARN string + EvaluationConfig []byte + InferenceConfig []byte + OutputDataS3URI string + ApplicationType string + ClientRequestToken string + CustomerEncryptionKeyID string + JobDescription string + JobTags map[string]string +} + +// EvaluationJob describes a model-evaluation job. +type EvaluationJob struct { + JobARN string + JobName string + JobType string + ApplicationType string + RoleARN string + EvaluationConfig []byte + InferenceConfig []byte + OutputDataS3URI string + JobDescription string + CustomerEncryptionKeyID string + Status string + FailureMessages []string + CreationTime string + LastModifiedTime string +} diff --git a/services/bedrock/driver/driver.go b/services/bedrock/driver/driver.go index 278643bc..2e80b256 100644 --- a/services/bedrock/driver/driver.go +++ b/services/bedrock/driver/driver.go @@ -166,6 +166,7 @@ type GuardrailConfig struct { KMSKeyID string ClientRequestToken string Tags map[string]string + GuardrailPolicies } // Guardrail describes a content guardrail. @@ -181,6 +182,7 @@ type Guardrail struct { KMSKeyARN string CreatedAt string UpdatedAt string + GuardrailPolicies } // ProvisionedThroughputConfig describes a provisioned-throughput purchase. @@ -231,6 +233,50 @@ type LoggingConfig struct { CloudWatch *CloudWatchLoggingConfig } +// Tag is a key/value label attached to a Bedrock resource, keyed by ARN. +type Tag struct { + Key string + Value string +} + +// CountTokensInput requests a token count for a would-be inference request. It +// carries exactly one of InvokeBody (a model-native InvokeModel payload) or the +// Converse messages/system, mirroring the CountTokens union. +type CountTokensInput struct { + ModelID string + InvokeBody []byte + Messages []Message + System []string +} + +// Guardrail source values for ApplyGuardrail. +const ( + GuardrailSourceInput = "INPUT" + GuardrailSourceOutput = "OUTPUT" +) + +// Guardrail intervention actions returned by ApplyGuardrail. +const ( + GuardrailActionNone = "NONE" + GuardrailActionIntervened = "GUARDRAIL_INTERVENED" +) + +// ApplyGuardrailInput evaluates content against a guardrail. +type ApplyGuardrailInput struct { + GuardrailIdentifier string + GuardrailVersion string + Source string + Content []string +} + +// ApplyGuardrailOutput is the result of evaluating content against a guardrail. +// The emulator never intervenes, so Action is always NONE and no policy units +// are consumed. +type ApplyGuardrailOutput struct { + Action string + Outputs []string +} + // Bedrock is the interface that foundation-model service implementations must // satisfy. It spans the control plane (model catalog, customization jobs, // custom models, guardrails, provisioned throughput, invocation logging) and @@ -249,12 +295,19 @@ type Bedrock interface { InvokeModel(ctx context.Context, in InvokeModelInput) (*InvokeModelResult, error) Converse(ctx context.Context, in ConverseInput) (*ConverseOutput, error) + CountTokens(ctx context.Context, in CountTokensInput) (int, error) + ApplyGuardrail(ctx context.Context, in ApplyGuardrailInput) (*ApplyGuardrailOutput, error) + + TagResource(ctx context.Context, resourceARN string, tags []Tag) error + UntagResource(ctx context.Context, resourceARN string, tagKeys []string) error + ListTagsForResource(ctx context.Context, resourceARN string) ([]Tag, error) CreateGuardrail(ctx context.Context, cfg GuardrailConfig) (*Guardrail, error) GetGuardrail(ctx context.Context, identifier, version string) (*Guardrail, error) - ListGuardrails(ctx context.Context) ([]Guardrail, error) + ListGuardrails(ctx context.Context, identifier string) ([]Guardrail, error) UpdateGuardrail(ctx context.Context, identifier string, cfg GuardrailConfig) (*Guardrail, error) - DeleteGuardrail(ctx context.Context, identifier string) error + CreateGuardrailVersion(ctx context.Context, identifier, description string) (guardrailID, version string, err error) + DeleteGuardrail(ctx context.Context, identifier, version string) error CreateProvisionedModelThroughput(ctx context.Context, cfg ProvisionedThroughputConfig) (*ProvisionedThroughput, error) GetProvisionedModelThroughput(ctx context.Context, identifier string) (*ProvisionedThroughput, error) @@ -264,4 +317,52 @@ type Bedrock interface { PutModelInvocationLoggingConfiguration(ctx context.Context, cfg LoggingConfig) error GetModelInvocationLoggingConfiguration(ctx context.Context) (*LoggingConfig, error) DeleteModelInvocationLoggingConfiguration(ctx context.Context) error + + StartAsyncInvoke(ctx context.Context, cfg StartAsyncInvokeConfig) (*AsyncInvoke, error) + GetAsyncInvoke(ctx context.Context, invocationARN string) (*AsyncInvoke, error) + ListAsyncInvokes(ctx context.Context) ([]AsyncInvoke, error) + + CreateModelImportJob(ctx context.Context, cfg ModelImportJobConfig) (*ModelImportJob, error) + GetModelImportJob(ctx context.Context, jobIdentifier string) (*ModelImportJob, error) + ListModelImportJobs(ctx context.Context) ([]ModelImportJob, error) + + CreateModelCopyJob(ctx context.Context, cfg ModelCopyJobConfig) (*ModelCopyJob, error) + GetModelCopyJob(ctx context.Context, jobARN string) (*ModelCopyJob, error) + ListModelCopyJobs(ctx context.Context) ([]ModelCopyJob, error) + + CreateEvaluationJob(ctx context.Context, cfg EvaluationJobConfig) (*EvaluationJob, error) + GetEvaluationJob(ctx context.Context, jobIdentifier string) (*EvaluationJob, error) + ListEvaluationJobs(ctx context.Context) ([]EvaluationJob, error) + StopEvaluationJob(ctx context.Context, jobIdentifier string) error + + CreateInferenceProfile(ctx context.Context, cfg InferenceProfileConfig) (*InferenceProfile, error) + GetInferenceProfile(ctx context.Context, identifier string) (*InferenceProfile, error) + ListInferenceProfiles(ctx context.Context) ([]InferenceProfile, error) + DeleteInferenceProfile(ctx context.Context, identifier string) error + + CreatePromptRouter(ctx context.Context, cfg PromptRouterConfig) (*PromptRouter, error) + GetPromptRouter(ctx context.Context, promptRouterARN string) (*PromptRouter, error) + ListPromptRouters(ctx context.Context) ([]PromptRouter, error) + DeletePromptRouter(ctx context.Context, promptRouterARN string) error + + CreateAutomatedReasoningPolicy(ctx context.Context, cfg AutomatedReasoningPolicyConfig) (*AutomatedReasoningPolicy, error) + GetAutomatedReasoningPolicy(ctx context.Context, policyARN string) (*AutomatedReasoningPolicy, error) + ListAutomatedReasoningPolicies(ctx context.Context) ([]AutomatedReasoningPolicy, error) + UpdateAutomatedReasoningPolicy( + ctx context.Context, policyARN string, upd AutomatedReasoningPolicyUpdate, + ) (*AutomatedReasoningPolicy, error) + DeleteAutomatedReasoningPolicy(ctx context.Context, policyARN string) error + + CreateMarketplaceModelEndpoint(ctx context.Context, cfg MarketplaceEndpointConfig) (*MarketplaceEndpoint, error) + GetMarketplaceModelEndpoint(ctx context.Context, endpointARN string) (*MarketplaceEndpoint, error) + ListMarketplaceModelEndpoints(ctx context.Context) ([]MarketplaceEndpoint, error) + UpdateMarketplaceModelEndpoint(ctx context.Context, endpointARN string, endpointConfig []byte) (*MarketplaceEndpoint, error) + DeleteMarketplaceModelEndpoint(ctx context.Context, endpointARN string) error + RegisterMarketplaceModelEndpoint(ctx context.Context, endpointIdentifier, modelSourceIdentifier string) (*MarketplaceEndpoint, error) + DeregisterMarketplaceModelEndpoint(ctx context.Context, endpointARN string) error + + CreateFoundationModelAgreement(ctx context.Context, modelID, offerToken string) (string, error) + DeleteFoundationModelAgreement(ctx context.Context, modelID string) error + ListFoundationModelAgreementOffers(ctx context.Context, modelID, offerType string) ([]FoundationModelOffer, error) + GetFoundationModelAvailability(ctx context.Context, modelID string) (*FoundationModelAvailability, error) } diff --git a/services/bedrock/driver/guardrail_policies.go b/services/bedrock/driver/guardrail_policies.go new file mode 100644 index 00000000..0873e7d9 --- /dev/null +++ b/services/bedrock/driver/guardrail_policies.go @@ -0,0 +1,110 @@ +package driver + +// Guardrail topic type values. +const ( + GuardrailTopicDeny = "DENY" +) + +// Guardrail content-filter strength values. +const ( + GuardrailStrengthNone = "NONE" + GuardrailStrengthLow = "LOW" + GuardrailStrengthMedium = "MEDIUM" + GuardrailStrengthHigh = "HIGH" +) + +// Guardrail PII / regex action values. +const ( + GuardrailPiiActionBlock = "BLOCK" + GuardrailPiiActionAnonymize = "ANONYMIZE" + GuardrailPiiActionNone = "NONE" +) + +// Guardrail contextual-grounding action values. +const ( + GuardrailGroundingActionBlock = "BLOCK" + GuardrailGroundingActionNone = "NONE" +) + +// GuardrailTopic is a single denied topic in a topic policy. +type GuardrailTopic struct { + Name string + Definition string + Examples []string + Type string +} + +// GuardrailTopicPolicy denies conversation on named topics. +type GuardrailTopicPolicy struct { + Topics []GuardrailTopic +} + +// GuardrailContentFilter is a single harmful-content filter. +type GuardrailContentFilter struct { + Type string + InputStrength string + OutputStrength string +} + +// GuardrailContentPolicy configures harmful-content filters. +type GuardrailContentPolicy struct { + Filters []GuardrailContentFilter +} + +// GuardrailWord is a single denied word or phrase. +type GuardrailWord struct { + Text string +} + +// GuardrailManagedWordList selects a managed word list (e.g. PROFANITY). +type GuardrailManagedWordList struct { + Type string +} + +// GuardrailWordPolicy configures denied words and managed word lists. +type GuardrailWordPolicy struct { + Words []GuardrailWord + ManagedWordLists []GuardrailManagedWordList +} + +// GuardrailPiiEntity configures handling of a PII entity type. +type GuardrailPiiEntity struct { + Type string + Action string +} + +// GuardrailRegex configures handling of a custom regular expression. +type GuardrailRegex struct { + Name string + Pattern string + Action string + Description string +} + +// GuardrailSensitiveInformationPolicy configures PII and regex handling. +type GuardrailSensitiveInformationPolicy struct { + PiiEntities []GuardrailPiiEntity + Regexes []GuardrailRegex +} + +// GuardrailContextualGroundingFilter is a single grounding/relevance filter. +type GuardrailContextualGroundingFilter struct { + Type string + Threshold float64 + Action string +} + +// GuardrailContextualGroundingPolicy configures grounding and relevance checks. +type GuardrailContextualGroundingPolicy struct { + Filters []GuardrailContextualGroundingFilter +} + +// GuardrailPolicies bundles the five configurable guardrail policies. It is +// embedded in both GuardrailConfig (request) and Guardrail (stored/response). +type GuardrailPolicies struct { + TopicPolicy *GuardrailTopicPolicy + ContentPolicy *GuardrailContentPolicy + WordPolicy *GuardrailWordPolicy + SensitiveInformationPolicy *GuardrailSensitiveInformationPolicy + ContextualGroundingPolicy *GuardrailContextualGroundingPolicy +} diff --git a/services/bedrock/driver/marketplace_agreements.go b/services/bedrock/driver/marketplace_agreements.go new file mode 100644 index 00000000..dda5ce23 --- /dev/null +++ b/services/bedrock/driver/marketplace_agreements.go @@ -0,0 +1,58 @@ +package driver + +// Marketplace model endpoint status values. Status is the registration status +// (types.Status); EndpointStatus mirrors the SageMaker endpoint state. +const ( + MarketplaceEndpointStatusRegistered = "REGISTERED" + MarketplaceEndpointStatusIncompatible = "INCOMPATIBLE_ENDPOINT" + MarketplaceEndpointStatusInService = "InService" +) + +// Foundation-model agreement / availability values. +const ( + AgreementStatusAvailable = "AVAILABLE" + AgreementStatusPending = "PENDING" + AuthorizationStatusAuthorized = "AUTHORIZED" + AuthorizationStatusNotAuthorized = "NOT_AUTHORIZED" + AvailabilityAvailable = "AVAILABLE" +) + +// MarketplaceEndpointConfig describes a marketplace model endpoint to create. +// EndpointConfig is the opaque endpointConfig union JSON carried through +// verbatim. +type MarketplaceEndpointConfig struct { + EndpointName string + ModelSourceIdentifier string + EndpointConfig []byte + AcceptEula bool + ClientRequestToken string + Tags map[string]string +} + +// MarketplaceEndpoint describes a deployed marketplace model endpoint. +type MarketplaceEndpoint struct { + EndpointARN string + ModelSourceIdentifier string + EndpointConfig []byte + EndpointStatus string + Status string + EndpointStatusMessage string + StatusMessage string + CreatedAt string + UpdatedAt string +} + +// FoundationModelOffer is a synthetic agreement offer for a foundation model. +type FoundationModelOffer struct { + OfferToken string + OfferID string +} + +// FoundationModelAvailability describes the agreement, authorization, +// entitlement, and region availability of a foundation model. +type FoundationModelAvailability struct { + AgreementStatus string + AuthorizationStatus string + EntitlementAvailability string + RegionAvailability string +} diff --git a/services/bedrock/driver/registries.go b/services/bedrock/driver/registries.go new file mode 100644 index 00000000..7c651a7f --- /dev/null +++ b/services/bedrock/driver/registries.go @@ -0,0 +1,104 @@ +package driver + +// Inference-profile status/type values. +const ( + InferenceProfileStatusActive = "ACTIVE" + InferenceProfileTypeApplication = "APPLICATION" + InferenceProfileTypeSystemDefined = "SYSTEM_DEFINED" +) + +// Prompt-router status/type values. +const ( + PromptRouterStatusAvailable = "AVAILABLE" + PromptRouterTypeCustom = "custom" + PromptRouterTypeDefault = "default" +) + +// AutomatedReasoningPolicyVersionDraft is the version assigned to newly created +// automated reasoning policies. +const AutomatedReasoningPolicyVersionDraft = "DRAFT" + +// InferenceProfileConfig describes an application inference profile to create. +// ModelSourceCopyFrom is the source model ARN from the modelSource copyFrom +// union member. +type InferenceProfileConfig struct { + Name string + ModelSourceCopyFrom string + ClientRequestToken string + Description string + Tags map[string]string +} + +// InferenceProfile describes an application inference profile. Models holds the +// tracked model ARNs. +type InferenceProfile struct { + ARN string + ID string + Name string + Models []string + Status string + Type string + Description string + CreatedAt string + UpdatedAt string +} + +// PromptRouterConfig describes a prompt router to create. Models and +// FallbackModelARN carry model ARNs; ResponseQualityDifference is the routing +// criterion. +type PromptRouterConfig struct { + Name string + Models []string + ResponseQualityDifference *float64 + FallbackModelARN string + ClientRequestToken string + Description string + Tags map[string]string +} + +// PromptRouter describes a prompt router. +type PromptRouter struct { + ARN string + Name string + Models []string + ResponseQualityDifference *float64 + FallbackModelARN string + Status string + Type string + Description string + CreatedAt string + UpdatedAt string +} + +// AutomatedReasoningPolicyConfig describes an automated reasoning policy to +// create. PolicyDefinition is an opaque JSON document carried through verbatim. +type AutomatedReasoningPolicyConfig struct { + Name string + ClientRequestToken string + Description string + KMSKeyID string + PolicyDefinition []byte + Tags map[string]string +} + +// AutomatedReasoningPolicyUpdate describes an update to an existing automated +// reasoning policy. +type AutomatedReasoningPolicyUpdate struct { + PolicyDefinition []byte + Description string + Name string +} + +// AutomatedReasoningPolicy describes an automated reasoning policy. +type AutomatedReasoningPolicy struct { + ARN string + ID string + Name string + Version string + DefinitionHash string + Description string + KMSKeyARN string + PolicyDefinition []byte + CreatedAt string + UpdatedAt string +} diff --git a/services/bedrock/marketplace_agreements.go b/services/bedrock/marketplace_agreements.go new file mode 100644 index 00000000..4408f449 --- /dev/null +++ b/services/bedrock/marketplace_agreements.go @@ -0,0 +1,146 @@ +package bedrock + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- Marketplace model endpoints --- + +// CreateMarketplaceModelEndpoint deploys a marketplace model endpoint. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *Bedrock) CreateMarketplaceModelEndpoint( + ctx context.Context, cfg driver.MarketplaceEndpointConfig, +) (*driver.MarketplaceEndpoint, error) { + out, err := b.do(ctx, "CreateMarketplaceModelEndpoint", cfg.EndpointName, func() (any, error) { + return b.driver.CreateMarketplaceModelEndpoint(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.MarketplaceEndpoint), nil +} + +// GetMarketplaceModelEndpoint retrieves a marketplace model endpoint by ARN. +func (b *Bedrock) GetMarketplaceModelEndpoint(ctx context.Context, endpointARN string) (*driver.MarketplaceEndpoint, error) { + out, err := b.do(ctx, "GetMarketplaceModelEndpoint", endpointARN, func() (any, error) { + return b.driver.GetMarketplaceModelEndpoint(ctx, endpointARN) + }) + if err != nil { + return nil, err + } + + return out.(*driver.MarketplaceEndpoint), nil +} + +// ListMarketplaceModelEndpoints lists all marketplace model endpoints. +func (b *Bedrock) ListMarketplaceModelEndpoints(ctx context.Context) ([]driver.MarketplaceEndpoint, error) { + out, err := b.do(ctx, "ListMarketplaceModelEndpoints", nil, func() (any, error) { + return b.driver.ListMarketplaceModelEndpoints(ctx) + }) + if err != nil { + return nil, err + } + + return out.([]driver.MarketplaceEndpoint), nil +} + +// UpdateMarketplaceModelEndpoint replaces an endpoint's configuration. +func (b *Bedrock) UpdateMarketplaceModelEndpoint( + ctx context.Context, endpointARN string, endpointConfig []byte, +) (*driver.MarketplaceEndpoint, error) { + out, err := b.do(ctx, "UpdateMarketplaceModelEndpoint", endpointARN, func() (any, error) { + return b.driver.UpdateMarketplaceModelEndpoint(ctx, endpointARN, endpointConfig) + }) + if err != nil { + return nil, err + } + + return out.(*driver.MarketplaceEndpoint), nil +} + +// DeleteMarketplaceModelEndpoint deletes a marketplace model endpoint by ARN. +func (b *Bedrock) DeleteMarketplaceModelEndpoint(ctx context.Context, endpointARN string) error { + _, err := b.do(ctx, "DeleteMarketplaceModelEndpoint", endpointARN, func() (any, error) { + return nil, b.driver.DeleteMarketplaceModelEndpoint(ctx, endpointARN) + }) + + return err +} + +// RegisterMarketplaceModelEndpoint registers an existing endpoint. +func (b *Bedrock) RegisterMarketplaceModelEndpoint( + ctx context.Context, endpointIdentifier, modelSourceIdentifier string, +) (*driver.MarketplaceEndpoint, error) { + out, err := b.do(ctx, "RegisterMarketplaceModelEndpoint", endpointIdentifier, func() (any, error) { + return b.driver.RegisterMarketplaceModelEndpoint(ctx, endpointIdentifier, modelSourceIdentifier) + }) + if err != nil { + return nil, err + } + + return out.(*driver.MarketplaceEndpoint), nil +} + +// DeregisterMarketplaceModelEndpoint deregisters an existing endpoint. +func (b *Bedrock) DeregisterMarketplaceModelEndpoint(ctx context.Context, endpointARN string) error { + _, err := b.do(ctx, "DeregisterMarketplaceModelEndpoint", endpointARN, func() (any, error) { + return nil, b.driver.DeregisterMarketplaceModelEndpoint(ctx, endpointARN) + }) + + return err +} + +// --- Foundation model agreements --- + +// CreateFoundationModelAgreement records an accepted agreement for a model. +func (b *Bedrock) CreateFoundationModelAgreement(ctx context.Context, modelID, offerToken string) (string, error) { + out, err := b.do(ctx, "CreateFoundationModelAgreement", modelID, func() (any, error) { + return b.driver.CreateFoundationModelAgreement(ctx, modelID, offerToken) + }) + if err != nil { + return "", err + } + + return out.(string), nil +} + +// DeleteFoundationModelAgreement removes an accepted agreement for a model. +func (b *Bedrock) DeleteFoundationModelAgreement(ctx context.Context, modelID string) error { + _, err := b.do(ctx, "DeleteFoundationModelAgreement", modelID, func() (any, error) { + return nil, b.driver.DeleteFoundationModelAgreement(ctx, modelID) + }) + + return err +} + +// ListFoundationModelAgreementOffers lists agreement offers for a model. +func (b *Bedrock) ListFoundationModelAgreementOffers( + ctx context.Context, modelID, offerType string, +) ([]driver.FoundationModelOffer, error) { + out, err := b.do(ctx, "ListFoundationModelAgreementOffers", modelID, func() (any, error) { + return b.driver.ListFoundationModelAgreementOffers(ctx, modelID, offerType) + }) + if err != nil { + return nil, err + } + + return out.([]driver.FoundationModelOffer), nil +} + +// GetFoundationModelAvailability reports availability for a model. +func (b *Bedrock) GetFoundationModelAvailability( + ctx context.Context, modelID string, +) (*driver.FoundationModelAvailability, error) { + out, err := b.do(ctx, "GetFoundationModelAvailability", modelID, func() (any, error) { + return b.driver.GetFoundationModelAvailability(ctx, modelID) + }) + if err != nil { + return nil, err + } + + return out.(*driver.FoundationModelAvailability), nil +} diff --git a/services/bedrock/registries.go b/services/bedrock/registries.go new file mode 100644 index 00000000..8ff4dd2f --- /dev/null +++ b/services/bedrock/registries.go @@ -0,0 +1,164 @@ +package bedrock + +import ( + "context" + + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// --- Inference profiles --- + +// CreateInferenceProfile creates an application inference profile. +func (b *Bedrock) CreateInferenceProfile(ctx context.Context, cfg driver.InferenceProfileConfig) (*driver.InferenceProfile, error) { + out, err := b.do(ctx, "CreateInferenceProfile", cfg.Name, func() (any, error) { + return b.driver.CreateInferenceProfile(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.InferenceProfile), nil +} + +// GetInferenceProfile retrieves an inference profile by ID or ARN. +func (b *Bedrock) GetInferenceProfile(ctx context.Context, identifier string) (*driver.InferenceProfile, error) { + out, err := b.do(ctx, "GetInferenceProfile", identifier, func() (any, error) { + return b.driver.GetInferenceProfile(ctx, identifier) + }) + if err != nil { + return nil, err + } + + return out.(*driver.InferenceProfile), nil +} + +// ListInferenceProfiles lists all inference profiles. +func (b *Bedrock) ListInferenceProfiles(ctx context.Context) ([]driver.InferenceProfile, error) { + out, err := b.do(ctx, "ListInferenceProfiles", nil, func() (any, error) { return b.driver.ListInferenceProfiles(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.InferenceProfile), nil +} + +// DeleteInferenceProfile deletes an inference profile by ID or ARN. +func (b *Bedrock) DeleteInferenceProfile(ctx context.Context, identifier string) error { + _, err := b.do(ctx, "DeleteInferenceProfile", identifier, func() (any, error) { + return nil, b.driver.DeleteInferenceProfile(ctx, identifier) + }) + + return err +} + +// --- Prompt routers --- + +// CreatePromptRouter creates a prompt router. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *Bedrock) CreatePromptRouter(ctx context.Context, cfg driver.PromptRouterConfig) (*driver.PromptRouter, error) { + out, err := b.do(ctx, "CreatePromptRouter", cfg.Name, func() (any, error) { + return b.driver.CreatePromptRouter(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.PromptRouter), nil +} + +// GetPromptRouter retrieves a prompt router by its ARN. +func (b *Bedrock) GetPromptRouter(ctx context.Context, promptRouterARN string) (*driver.PromptRouter, error) { + out, err := b.do(ctx, "GetPromptRouter", promptRouterARN, func() (any, error) { + return b.driver.GetPromptRouter(ctx, promptRouterARN) + }) + if err != nil { + return nil, err + } + + return out.(*driver.PromptRouter), nil +} + +// ListPromptRouters lists all prompt routers. +func (b *Bedrock) ListPromptRouters(ctx context.Context) ([]driver.PromptRouter, error) { + out, err := b.do(ctx, "ListPromptRouters", nil, func() (any, error) { return b.driver.ListPromptRouters(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.PromptRouter), nil +} + +// DeletePromptRouter deletes a prompt router by its ARN. +func (b *Bedrock) DeletePromptRouter(ctx context.Context, promptRouterARN string) error { + _, err := b.do(ctx, "DeletePromptRouter", promptRouterARN, func() (any, error) { + return nil, b.driver.DeletePromptRouter(ctx, promptRouterARN) + }) + + return err +} + +// --- Automated reasoning policies --- + +// CreateAutomatedReasoningPolicy creates an automated reasoning policy. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *Bedrock) CreateAutomatedReasoningPolicy( + ctx context.Context, cfg driver.AutomatedReasoningPolicyConfig, +) (*driver.AutomatedReasoningPolicy, error) { + out, err := b.do(ctx, "CreateAutomatedReasoningPolicy", cfg.Name, func() (any, error) { + return b.driver.CreateAutomatedReasoningPolicy(ctx, cfg) + }) + if err != nil { + return nil, err + } + + return out.(*driver.AutomatedReasoningPolicy), nil +} + +// GetAutomatedReasoningPolicy retrieves an automated reasoning policy by its ARN. +func (b *Bedrock) GetAutomatedReasoningPolicy(ctx context.Context, policyARN string) (*driver.AutomatedReasoningPolicy, error) { + out, err := b.do(ctx, "GetAutomatedReasoningPolicy", policyARN, func() (any, error) { + return b.driver.GetAutomatedReasoningPolicy(ctx, policyARN) + }) + if err != nil { + return nil, err + } + + return out.(*driver.AutomatedReasoningPolicy), nil +} + +// ListAutomatedReasoningPolicies lists all automated reasoning policies. +func (b *Bedrock) ListAutomatedReasoningPolicies(ctx context.Context) ([]driver.AutomatedReasoningPolicy, error) { + out, err := b.do(ctx, "ListAutomatedReasoningPolicies", nil, func() (any, error) { + return b.driver.ListAutomatedReasoningPolicies(ctx) + }) + if err != nil { + return nil, err + } + + return out.([]driver.AutomatedReasoningPolicy), nil +} + +// UpdateAutomatedReasoningPolicy updates an automated reasoning policy. +func (b *Bedrock) UpdateAutomatedReasoningPolicy( + ctx context.Context, policyARN string, upd driver.AutomatedReasoningPolicyUpdate, +) (*driver.AutomatedReasoningPolicy, error) { + out, err := b.do(ctx, "UpdateAutomatedReasoningPolicy", policyARN, func() (any, error) { + return b.driver.UpdateAutomatedReasoningPolicy(ctx, policyARN, upd) + }) + if err != nil { + return nil, err + } + + return out.(*driver.AutomatedReasoningPolicy), nil +} + +// DeleteAutomatedReasoningPolicy deletes an automated reasoning policy by its ARN. +func (b *Bedrock) DeleteAutomatedReasoningPolicy(ctx context.Context, policyARN string) error { + _, err := b.do(ctx, "DeleteAutomatedReasoningPolicy", policyARN, func() (any, error) { + return nil, b.driver.DeleteAutomatedReasoningPolicy(ctx, policyARN) + }) + + return err +} diff --git a/services/bedrockagent/bedrockagent.go b/services/bedrockagent/bedrockagent.go new file mode 100644 index 00000000..55af95f2 --- /dev/null +++ b/services/bedrockagent/bedrockagent.go @@ -0,0 +1,413 @@ +// Package bedrockagent provides a portable Bedrock Agent authoring API with +// cross-cutting concerns. It wraps a driver.BedrockAgent with recording, +// metrics, rate limiting, error injection, and latency simulation. +package bedrockagent + +import ( + "context" + "time" + + "github.com/stackshy/cloudemu/v2/features/inject" + "github.com/stackshy/cloudemu/v2/features/metrics" + "github.com/stackshy/cloudemu/v2/features/ratelimit" + "github.com/stackshy/cloudemu/v2/features/recorder" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +const service = "bedrockagent" + +// BedrockAgent is the portable Bedrock Agent type wrapping a driver with +// cross-cutting concerns. +type BedrockAgent struct { + driver driver.BedrockAgent + recorder *recorder.Recorder + metrics *metrics.Collector + limiter *ratelimit.Limiter + injector *inject.Injector + latency time.Duration +} + +// NewBedrockAgent creates a new portable BedrockAgent wrapping the given driver. +func NewBedrockAgent(d driver.BedrockAgent, opts ...Option) *BedrockAgent { + b := &BedrockAgent{driver: d} + for _, opt := range opts { + opt(b) + } + + return b +} + +// Option configures a portable BedrockAgent. +type Option func(*BedrockAgent) + +// WithRecorder sets the recorder. +func WithRecorder(r *recorder.Recorder) Option { return func(b *BedrockAgent) { b.recorder = r } } + +// WithMetrics sets the metrics collector. +func WithMetrics(m *metrics.Collector) Option { return func(b *BedrockAgent) { b.metrics = m } } + +// WithRateLimiter sets the rate limiter. +func WithRateLimiter(l *ratelimit.Limiter) Option { return func(b *BedrockAgent) { b.limiter = l } } + +// WithErrorInjection sets the error injector. +func WithErrorInjection(i *inject.Injector) Option { return func(b *BedrockAgent) { b.injector = i } } + +// WithLatency sets simulated latency. +func WithLatency(d time.Duration) Option { return func(b *BedrockAgent) { b.latency = d } } + +func (b *BedrockAgent) do(_ context.Context, op string, input any, fn func() (any, error)) (any, error) { + start := time.Now() + + if b.injector != nil { + if err := b.injector.Check(service, op); err != nil { + b.rec(op, input, nil, err, time.Since(start)) + return nil, err + } + } + + if b.limiter != nil { + if err := b.limiter.Allow(); err != nil { + b.rec(op, input, nil, err, time.Since(start)) + return nil, err + } + } + + if b.latency > 0 { + time.Sleep(b.latency) + } + + out, err := fn() + dur := time.Since(start) + + if b.metrics != nil { + labels := map[string]string{"service": service, "operation": op} + b.metrics.Counter("calls_total", 1, labels) + b.metrics.Histogram("call_duration", dur, labels) + + if err != nil { + b.metrics.Counter("errors_total", 1, labels) + } + } + + b.rec(op, input, out, err, dur) + + return out, err +} + +func (b *BedrockAgent) rec(op string, input, output any, err error, dur time.Duration) { + if b.recorder != nil { + b.recorder.Record(service, op, input, output, err, dur) + } +} + +// CreateAgent creates a new agent. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) CreateAgent(ctx context.Context, cfg driver.AgentConfig) (*driver.Agent, error) { + out, err := b.do(ctx, "CreateAgent", cfg, func() (any, error) { return b.driver.CreateAgent(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.Agent), nil +} + +// GetAgent retrieves an agent by ID. +func (b *BedrockAgent) GetAgent(ctx context.Context, agentID string) (*driver.Agent, error) { + out, err := b.do(ctx, "GetAgent", agentID, func() (any, error) { return b.driver.GetAgent(ctx, agentID) }) + if err != nil { + return nil, err + } + + return out.(*driver.Agent), nil +} + +// ListAgents lists all agents. +func (b *BedrockAgent) ListAgents(ctx context.Context) ([]driver.Agent, error) { + out, err := b.do(ctx, "ListAgents", nil, func() (any, error) { return b.driver.ListAgents(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.Agent), nil +} + +// UpdateAgent updates an agent's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) UpdateAgent(ctx context.Context, agentID string, cfg driver.AgentConfig) (*driver.Agent, error) { + out, err := b.do(ctx, "UpdateAgent", agentID, func() (any, error) { return b.driver.UpdateAgent(ctx, agentID, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.Agent), nil +} + +// DeleteAgent deletes an agent and returns its terminal status. +func (b *BedrockAgent) DeleteAgent(ctx context.Context, agentID string) (string, error) { + out, err := b.do(ctx, "DeleteAgent", agentID, func() (any, error) { return b.driver.DeleteAgent(ctx, agentID) }) + if err != nil { + return "", err + } + + return out.(string), nil +} + +// PrepareAgent prepares an agent, transitioning it to PREPARED. +func (b *BedrockAgent) PrepareAgent(ctx context.Context, agentID string) (*driver.Agent, error) { + out, err := b.do(ctx, "PrepareAgent", agentID, func() (any, error) { return b.driver.PrepareAgent(ctx, agentID) }) + if err != nil { + return nil, err + } + + return out.(*driver.Agent), nil +} + +// CreateAgentAlias creates an alias of an agent. +func (b *BedrockAgent) CreateAgentAlias(ctx context.Context, cfg driver.AgentAliasConfig) (*driver.AgentAlias, error) { + out, err := b.do(ctx, "CreateAgentAlias", cfg, func() (any, error) { return b.driver.CreateAgentAlias(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.AgentAlias), nil +} + +// CreateKnowledgeBase creates a knowledge base. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) CreateKnowledgeBase(ctx context.Context, cfg driver.KnowledgeBaseConfig) (*driver.KnowledgeBase, error) { + out, err := b.do(ctx, "CreateKnowledgeBase", cfg, func() (any, error) { return b.driver.CreateKnowledgeBase(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.KnowledgeBase), nil +} + +// GetKnowledgeBase retrieves a knowledge base by ID. +func (b *BedrockAgent) GetKnowledgeBase(ctx context.Context, id string) (*driver.KnowledgeBase, error) { + out, err := b.do(ctx, "GetKnowledgeBase", id, func() (any, error) { return b.driver.GetKnowledgeBase(ctx, id) }) + if err != nil { + return nil, err + } + + return out.(*driver.KnowledgeBase), nil +} + +// ListKnowledgeBases lists all knowledge bases. +func (b *BedrockAgent) ListKnowledgeBases(ctx context.Context) ([]driver.KnowledgeBase, error) { + out, err := b.do(ctx, "ListKnowledgeBases", nil, func() (any, error) { return b.driver.ListKnowledgeBases(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.KnowledgeBase), nil +} + +// UpdateKnowledgeBase updates a knowledge base's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) UpdateKnowledgeBase(ctx context.Context, id string, cfg driver.KnowledgeBaseConfig) (*driver.KnowledgeBase, error) { + out, err := b.do(ctx, "UpdateKnowledgeBase", id, func() (any, error) { return b.driver.UpdateKnowledgeBase(ctx, id, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.KnowledgeBase), nil +} + +// DeleteKnowledgeBase deletes a knowledge base and returns its terminal status. +func (b *BedrockAgent) DeleteKnowledgeBase(ctx context.Context, id string) (string, error) { + out, err := b.do(ctx, "DeleteKnowledgeBase", id, func() (any, error) { return b.driver.DeleteKnowledgeBase(ctx, id) }) + if err != nil { + return "", err + } + + return out.(string), nil +} + +// CreateDataSource creates a data source under a knowledge base. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) CreateDataSource(ctx context.Context, cfg driver.DataSourceConfig) (*driver.DataSource, error) { + out, err := b.do(ctx, "CreateDataSource", cfg, func() (any, error) { return b.driver.CreateDataSource(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.DataSource), nil +} + +// GetDataSource retrieves a data source by knowledge-base and data-source ID. +func (b *BedrockAgent) GetDataSource(ctx context.Context, kbID, dsID string) (*driver.DataSource, error) { + out, err := b.do(ctx, "GetDataSource", dsID, func() (any, error) { return b.driver.GetDataSource(ctx, kbID, dsID) }) + if err != nil { + return nil, err + } + + return out.(*driver.DataSource), nil +} + +// ListDataSources lists all data sources under a knowledge base. +func (b *BedrockAgent) ListDataSources(ctx context.Context, kbID string) ([]driver.DataSource, error) { + out, err := b.do(ctx, "ListDataSources", kbID, func() (any, error) { return b.driver.ListDataSources(ctx, kbID) }) + if err != nil { + return nil, err + } + + return out.([]driver.DataSource), nil +} + +// UpdateDataSource updates a data source's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) UpdateDataSource(ctx context.Context, cfg driver.DataSourceConfig, dsID string) (*driver.DataSource, error) { + out, err := b.do(ctx, "UpdateDataSource", dsID, func() (any, error) { return b.driver.UpdateDataSource(ctx, cfg, dsID) }) + if err != nil { + return nil, err + } + + return out.(*driver.DataSource), nil +} + +// DeleteDataSource deletes a data source and returns its terminal status. +func (b *BedrockAgent) DeleteDataSource(ctx context.Context, kbID, dsID string) (string, error) { + out, err := b.do(ctx, "DeleteDataSource", dsID, func() (any, error) { return b.driver.DeleteDataSource(ctx, kbID, dsID) }) + if err != nil { + return "", err + } + + return out.(string), nil +} + +// StartIngestionJob starts an ingestion job for a data source. +func (b *BedrockAgent) StartIngestionJob(ctx context.Context, kbID, dsID, description string) (*driver.IngestionJob, error) { + out, err := b.do(ctx, "StartIngestionJob", dsID, func() (any, error) { + return b.driver.StartIngestionJob(ctx, kbID, dsID, description) + }) + if err != nil { + return nil, err + } + + return out.(*driver.IngestionJob), nil +} + +// CreateFlow creates a flow. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) CreateFlow(ctx context.Context, cfg driver.FlowConfig) (*driver.Flow, error) { + out, err := b.do(ctx, "CreateFlow", cfg, func() (any, error) { return b.driver.CreateFlow(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.Flow), nil +} + +// GetFlow retrieves a flow by identifier. +func (b *BedrockAgent) GetFlow(ctx context.Context, id string) (*driver.Flow, error) { + out, err := b.do(ctx, "GetFlow", id, func() (any, error) { return b.driver.GetFlow(ctx, id) }) + if err != nil { + return nil, err + } + + return out.(*driver.Flow), nil +} + +// ListFlows lists all flows. +func (b *BedrockAgent) ListFlows(ctx context.Context) ([]driver.Flow, error) { + out, err := b.do(ctx, "ListFlows", nil, func() (any, error) { return b.driver.ListFlows(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.Flow), nil +} + +// UpdateFlow updates a flow's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) UpdateFlow(ctx context.Context, id string, cfg driver.FlowConfig) (*driver.Flow, error) { + out, err := b.do(ctx, "UpdateFlow", id, func() (any, error) { return b.driver.UpdateFlow(ctx, id, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.Flow), nil +} + +// DeleteFlow deletes a flow and returns its identifier. +func (b *BedrockAgent) DeleteFlow(ctx context.Context, id string) (string, error) { + out, err := b.do(ctx, "DeleteFlow", id, func() (any, error) { return b.driver.DeleteFlow(ctx, id) }) + if err != nil { + return "", err + } + + return out.(string), nil +} + +// PrepareFlow prepares a flow, transitioning it to Prepared. +func (b *BedrockAgent) PrepareFlow(ctx context.Context, id string) (*driver.Flow, error) { + out, err := b.do(ctx, "PrepareFlow", id, func() (any, error) { return b.driver.PrepareFlow(ctx, id) }) + if err != nil { + return nil, err + } + + return out.(*driver.Flow), nil +} + +// CreatePrompt creates a prompt. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) CreatePrompt(ctx context.Context, cfg driver.PromptConfig) (*driver.Prompt, error) { + out, err := b.do(ctx, "CreatePrompt", cfg, func() (any, error) { return b.driver.CreatePrompt(ctx, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.Prompt), nil +} + +// GetPrompt retrieves a prompt by identifier. +func (b *BedrockAgent) GetPrompt(ctx context.Context, id string) (*driver.Prompt, error) { + out, err := b.do(ctx, "GetPrompt", id, func() (any, error) { return b.driver.GetPrompt(ctx, id) }) + if err != nil { + return nil, err + } + + return out.(*driver.Prompt), nil +} + +// ListPrompts lists all prompts. +func (b *BedrockAgent) ListPrompts(ctx context.Context) ([]driver.Prompt, error) { + out, err := b.do(ctx, "ListPrompts", nil, func() (any, error) { return b.driver.ListPrompts(ctx) }) + if err != nil { + return nil, err + } + + return out.([]driver.Prompt), nil +} + +// UpdatePrompt updates a prompt's mutable fields. +// +//nolint:gocritic // cfg matches the driver interface signature; copied once on entry. +func (b *BedrockAgent) UpdatePrompt(ctx context.Context, id string, cfg driver.PromptConfig) (*driver.Prompt, error) { + out, err := b.do(ctx, "UpdatePrompt", id, func() (any, error) { return b.driver.UpdatePrompt(ctx, id, cfg) }) + if err != nil { + return nil, err + } + + return out.(*driver.Prompt), nil +} + +// DeletePrompt deletes a prompt and returns its identifier. +func (b *BedrockAgent) DeletePrompt(ctx context.Context, id string) (string, error) { + out, err := b.do(ctx, "DeletePrompt", id, func() (any, error) { return b.driver.DeletePrompt(ctx, id) }) + if err != nil { + return "", err + } + + return out.(string), nil +} diff --git a/services/bedrockagent/driver/driver.go b/services/bedrockagent/driver/driver.go new file mode 100644 index 00000000..def26de8 --- /dev/null +++ b/services/bedrockagent/driver/driver.go @@ -0,0 +1,238 @@ +// Package driver defines the interface for Bedrock Agent-style services: the +// authoring control plane for agents, knowledge bases, data sources, ingestion +// jobs, flows, and prompts. +package driver + +import ( + "context" + "encoding/json" +) + +// Agent lifecycle status values (UPPER_SNAKE, per AgentStatus). +const ( + AgentNotPrepared = "NOT_PREPARED" + AgentPreparing = "PREPARING" + AgentPrepared = "PREPARED" +) + +// AgentAlias status values. +const ( + AgentAliasPrepared = "PREPARED" +) + +// KnowledgeBase status values. +const ( + KnowledgeBaseActive = "ACTIVE" +) + +// DataSource status values. +const ( + DataSourceAvailable = "AVAILABLE" +) + +// IngestionJob status values. +const ( + IngestionJobComplete = "COMPLETE" +) + +// Flow status values (PascalCase, per FlowStatus). +const ( + FlowNotPrepared = "NotPrepared" + FlowPreparing = "Preparing" + FlowPrepared = "Prepared" +) + +// DraftVersion is the working version assigned to freshly created resources. +const DraftVersion = "DRAFT" + +// AgentConfig describes an agent to create. +type AgentConfig struct { + Name string + ResourceRoleArn string + FoundationModel string + Instruction string + Description string + IdleSessionTTLInSeconds int32 + ClientToken string + Tags map[string]string +} + +// Agent describes a Bedrock agent. +type Agent struct { + ID string + ARN string + Name string + ResourceRoleArn string + FoundationModel string + Instruction string + Description string + Status string + Version string + IdleSessionTTLInSeconds int32 + CreatedAt string + UpdatedAt string + PreparedAt string +} + +// AgentAliasConfig describes an agent alias to create. +type AgentAliasConfig struct { + AgentID string + Name string + Description string +} + +// AgentAlias describes an alias of an agent. +type AgentAlias struct { + ID string + ARN string + AgentID string + Name string + Description string + Status string + CreatedAt string + UpdatedAt string +} + +// KnowledgeBaseConfig describes a knowledge base to create. +type KnowledgeBaseConfig struct { + Name string + RoleArn string + Description string + KnowledgeBaseConfiguration json.RawMessage + StorageConfiguration json.RawMessage + Tags map[string]string +} + +// KnowledgeBase describes a Bedrock knowledge base. +type KnowledgeBase struct { + ID string + ARN string + Name string + RoleArn string + Description string + Status string + KnowledgeBaseConfiguration json.RawMessage + StorageConfiguration json.RawMessage + CreatedAt string + UpdatedAt string +} + +// DataSourceConfig describes a data source to create. +type DataSourceConfig struct { + KnowledgeBaseID string + Name string + Description string + DataDeletionPolicy string + DataSourceConfiguration json.RawMessage + VectorIngestionConfiguration json.RawMessage +} + +// DataSource describes a knowledge-base data source. +type DataSource struct { + ID string + KnowledgeBaseID string + Name string + Description string + Status string + DataDeletionPolicy string + DataSourceConfiguration json.RawMessage + CreatedAt string + UpdatedAt string +} + +// IngestionJob describes a data-source ingestion job. +type IngestionJob struct { + ID string + KnowledgeBaseID string + DataSourceID string + Description string + Status string + StartedAt string + UpdatedAt string +} + +// FlowConfig describes a flow to create. +type FlowConfig struct { + Name string + ExecutionRoleArn string + Description string + CustomerEncryptionKeyArn string + Definition json.RawMessage +} + +// Flow describes a Bedrock flow. +type Flow struct { + ID string + ARN string + Name string + ExecutionRoleArn string + Description string + Status string + Version string + CustomerEncryptionKeyArn string + Definition json.RawMessage + CreatedAt string + UpdatedAt string +} + +// PromptConfig describes a prompt to create. +type PromptConfig struct { + Name string + Description string + DefaultVariant string + CustomerEncryptionKeyArn string + Variants json.RawMessage +} + +// Prompt describes a Bedrock prompt. +type Prompt struct { + ID string + ARN string + Name string + Description string + Version string + DefaultVariant string + CustomerEncryptionKeyArn string + Variants json.RawMessage + CreatedAt string + UpdatedAt string +} + +// BedrockAgent is the interface that Bedrock Agent authoring implementations +// must satisfy: agents (and aliases), knowledge bases, data sources, ingestion +// jobs, flows, and prompts. +type BedrockAgent interface { + CreateAgent(ctx context.Context, cfg AgentConfig) (*Agent, error) + GetAgent(ctx context.Context, agentID string) (*Agent, error) + ListAgents(ctx context.Context) ([]Agent, error) + UpdateAgent(ctx context.Context, agentID string, cfg AgentConfig) (*Agent, error) + DeleteAgent(ctx context.Context, agentID string) (string, error) + PrepareAgent(ctx context.Context, agentID string) (*Agent, error) + CreateAgentAlias(ctx context.Context, cfg AgentAliasConfig) (*AgentAlias, error) + + CreateKnowledgeBase(ctx context.Context, cfg KnowledgeBaseConfig) (*KnowledgeBase, error) + GetKnowledgeBase(ctx context.Context, id string) (*KnowledgeBase, error) + ListKnowledgeBases(ctx context.Context) ([]KnowledgeBase, error) + UpdateKnowledgeBase(ctx context.Context, id string, cfg KnowledgeBaseConfig) (*KnowledgeBase, error) + DeleteKnowledgeBase(ctx context.Context, id string) (string, error) + + CreateDataSource(ctx context.Context, cfg DataSourceConfig) (*DataSource, error) + GetDataSource(ctx context.Context, kbID, dsID string) (*DataSource, error) + ListDataSources(ctx context.Context, kbID string) ([]DataSource, error) + UpdateDataSource(ctx context.Context, cfg DataSourceConfig, dsID string) (*DataSource, error) + DeleteDataSource(ctx context.Context, kbID, dsID string) (string, error) + StartIngestionJob(ctx context.Context, kbID, dsID, description string) (*IngestionJob, error) + + CreateFlow(ctx context.Context, cfg FlowConfig) (*Flow, error) + GetFlow(ctx context.Context, id string) (*Flow, error) + ListFlows(ctx context.Context) ([]Flow, error) + UpdateFlow(ctx context.Context, id string, cfg FlowConfig) (*Flow, error) + DeleteFlow(ctx context.Context, id string) (string, error) + PrepareFlow(ctx context.Context, id string) (*Flow, error) + + CreatePrompt(ctx context.Context, cfg PromptConfig) (*Prompt, error) + GetPrompt(ctx context.Context, id string) (*Prompt, error) + ListPrompts(ctx context.Context) ([]Prompt, error) + UpdatePrompt(ctx context.Context, id string, cfg PromptConfig) (*Prompt, error) + DeletePrompt(ctx context.Context, id string) (string, error) +} diff --git a/services/bedrockagentruntime/bedrockagentruntime.go b/services/bedrockagentruntime/bedrockagentruntime.go new file mode 100644 index 00000000..c8f095cf --- /dev/null +++ b/services/bedrockagentruntime/bedrockagentruntime.go @@ -0,0 +1,141 @@ +// Package bedrockagentruntime provides a portable Bedrock Agent runtime API +// with cross-cutting concerns. It wraps a driver.BedrockAgentRuntime with +// recording, metrics, rate limiting, error injection, and latency simulation. +package bedrockagentruntime + +import ( + "context" + "time" + + "github.com/stackshy/cloudemu/v2/features/inject" + "github.com/stackshy/cloudemu/v2/features/metrics" + "github.com/stackshy/cloudemu/v2/features/ratelimit" + "github.com/stackshy/cloudemu/v2/features/recorder" + "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +// BedrockAgentRuntime is the portable Bedrock Agent runtime type wrapping a +// driver with cross-cutting concerns. +type BedrockAgentRuntime struct { + driver driver.BedrockAgentRuntime + recorder *recorder.Recorder + metrics *metrics.Collector + limiter *ratelimit.Limiter + injector *inject.Injector + latency time.Duration +} + +// NewBedrockAgentRuntime creates a new portable BedrockAgentRuntime wrapping the +// given driver. +func NewBedrockAgentRuntime(d driver.BedrockAgentRuntime, opts ...Option) *BedrockAgentRuntime { + b := &BedrockAgentRuntime{driver: d} + for _, opt := range opts { + opt(b) + } + + return b +} + +// Option configures a portable BedrockAgentRuntime. +type Option func(*BedrockAgentRuntime) + +// WithRecorder sets the recorder. +func WithRecorder(r *recorder.Recorder) Option { + return func(b *BedrockAgentRuntime) { b.recorder = r } +} + +// WithMetrics sets the metrics collector. +func WithMetrics(m *metrics.Collector) Option { return func(b *BedrockAgentRuntime) { b.metrics = m } } + +// WithRateLimiter sets the rate limiter. +func WithRateLimiter(l *ratelimit.Limiter) Option { + return func(b *BedrockAgentRuntime) { b.limiter = l } +} + +// WithErrorInjection sets the error injector. +func WithErrorInjection(i *inject.Injector) Option { + return func(b *BedrockAgentRuntime) { b.injector = i } +} + +// WithLatency sets simulated latency. +func WithLatency(d time.Duration) Option { return func(b *BedrockAgentRuntime) { b.latency = d } } + +func (b *BedrockAgentRuntime) do(_ context.Context, op string, input any, fn func() (any, error)) (any, error) { + start := time.Now() + + if b.injector != nil { + if err := b.injector.Check("bedrockagentruntime", op); err != nil { + b.rec(op, input, nil, err, time.Since(start)) + return nil, err + } + } + + if b.limiter != nil { + if err := b.limiter.Allow(); err != nil { + b.rec(op, input, nil, err, time.Since(start)) + return nil, err + } + } + + if b.latency > 0 { + time.Sleep(b.latency) + } + + out, err := fn() + dur := time.Since(start) + + if b.metrics != nil { + labels := map[string]string{"service": "bedrockagentruntime", "operation": op} + b.metrics.Counter("calls_total", 1, labels) + b.metrics.Histogram("call_duration", dur, labels) + + if err != nil { + b.metrics.Counter("errors_total", 1, labels) + } + } + + b.rec(op, input, out, err, dur) + + return out, err +} + +func (b *BedrockAgentRuntime) rec(op string, input, output any, err error, dur time.Duration) { + if b.recorder != nil { + b.recorder.Record("bedrockagentruntime", op, input, output, err, dur) + } +} + +// InvokeAgent runs (emulated) agent inference and returns the assembled +// completion. +func (b *BedrockAgentRuntime) InvokeAgent(ctx context.Context, in driver.InvokeAgentInput) (*driver.InvokeAgentResult, error) { + out, err := b.do(ctx, "InvokeAgent", in.AgentID, func() (any, error) { return b.driver.InvokeAgent(ctx, in) }) + if err != nil { + return nil, err + } + + return out.(*driver.InvokeAgentResult), nil +} + +// Retrieve queries a knowledge base and returns matching chunks. +func (b *BedrockAgentRuntime) Retrieve(ctx context.Context, in driver.RetrieveInput) (*driver.RetrieveResult, error) { + out, err := b.do(ctx, "Retrieve", in.KnowledgeBaseID, func() (any, error) { return b.driver.Retrieve(ctx, in) }) + if err != nil { + return nil, err + } + + return out.(*driver.RetrieveResult), nil +} + +// RetrieveAndGenerate queries a knowledge base and generates an answer. +func (b *BedrockAgentRuntime) RetrieveAndGenerate( + ctx context.Context, in driver.RetrieveAndGenerateInput, +) (*driver.RetrieveAndGenerateResult, error) { + out, err := b.do(ctx, "RetrieveAndGenerate", in.SessionID, func() (any, error) { + return b.driver.RetrieveAndGenerate(ctx, in) + }) + if err != nil { + return nil, err + } + + return out.(*driver.RetrieveAndGenerateResult), nil +} diff --git a/services/bedrockagentruntime/bedrockagentruntime_test.go b/services/bedrockagentruntime/bedrockagentruntime_test.go new file mode 100644 index 00000000..71ef8215 --- /dev/null +++ b/services/bedrockagentruntime/bedrockagentruntime_test.go @@ -0,0 +1,66 @@ +package bedrockagentruntime + +import ( + "context" + "strings" + "testing" + "time" + + "github.com/stackshy/cloudemu/v2/config" + provider "github.com/stackshy/cloudemu/v2/providers/aws/bedrockagentruntime" + "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" +) + +func newService() *BedrockAgentRuntime { + fc := config.NewFakeClock(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)) + opts := config.NewOptions(config.WithClock(fc), config.WithRegion("us-east-1")) + + return NewBedrockAgentRuntime(provider.New(opts)) +} + +func TestServiceInvokeAgent(t *testing.T) { + svc := newService() + + out, err := svc.InvokeAgent(context.Background(), driver.InvokeAgentInput{ + AgentID: "AGENT1", + AgentAliasID: "ALIAS1", + SessionID: "s1", + InputText: "hi", + }) + if err != nil { + t.Fatalf("InvokeAgent: %v", err) + } + + if !strings.Contains(out.Completion, "hi") || out.SessionID != "s1" { + t.Fatalf("unexpected result: %+v", out) + } +} + +func TestServiceRetrieve(t *testing.T) { + svc := newService() + + out, err := svc.Retrieve(context.Background(), driver.RetrieveInput{ + KnowledgeBaseID: "kb1", + QueryText: "q", + }) + if err != nil { + t.Fatalf("Retrieve: %v", err) + } + + if len(out.Results) == 0 { + t.Fatal("expected results") + } +} + +func TestServiceRetrieveAndGenerate(t *testing.T) { + svc := newService() + + out, err := svc.RetrieveAndGenerate(context.Background(), driver.RetrieveAndGenerateInput{InputText: "q"}) + if err != nil { + t.Fatalf("RetrieveAndGenerate: %v", err) + } + + if out.Text == "" || out.SessionID == "" { + t.Fatalf("unexpected result: %+v", out) + } +} diff --git a/services/bedrockagentruntime/driver/driver.go b/services/bedrockagentruntime/driver/driver.go new file mode 100644 index 00000000..689ed320 --- /dev/null +++ b/services/bedrockagentruntime/driver/driver.go @@ -0,0 +1,74 @@ +// Package driver defines the interface for the Bedrock Agent runtime +// (bedrock-agent-runtime) data-plane: invoking agents, retrieving from +// knowledge bases, and retrieve-and-generate. +// +// InvokeAgent's real API streams its answer as an eventstream of chunk events. +// The driver models the response as the fully-assembled completion text plus a +// session id; the server layer is responsible for splitting that text into +// eventstream chunk frames on the wire. +package driver + +import "context" + +// InvokeAgentInput is the request for the InvokeAgent operation. AgentID, +// AgentAliasID, and SessionID are path parameters; InputText is the prompt. +type InvokeAgentInput struct { + AgentID string + AgentAliasID string + SessionID string + InputText string + EnableTrace bool + EndSession bool +} + +// InvokeAgentResult carries the assembled agent completion. ContentType is the +// media type of the completion payload (application/json for text answers). +type InvokeAgentResult struct { + Completion string + SessionID string + ContentType string +} + +// RetrieveInput is the request for the Retrieve operation. KnowledgeBaseID is a +// path parameter; QueryText is the retrieval query text (required). +type RetrieveInput struct { + KnowledgeBaseID string + QueryText string + NextToken string +} + +// RetrievalResult is a single chunk returned from a knowledge-base query. +type RetrievalResult struct { + Text string + Score float64 + LocationURI string +} + +// RetrieveResult is the response from the Retrieve operation. +type RetrieveResult struct { + Results []RetrievalResult + NextToken string +} + +// RetrieveAndGenerateInput is the request for the RetrieveAndGenerate +// operation. InputText is the query (required); SessionID is optional and +// continues an existing session when supplied. +type RetrieveAndGenerateInput struct { + InputText string + SessionID string +} + +// RetrieveAndGenerateResult is the response from the RetrieveAndGenerate +// operation. +type RetrieveAndGenerateResult struct { + Text string + SessionID string +} + +// BedrockAgentRuntime is the interface that Bedrock Agent runtime +// implementations must satisfy. +type BedrockAgentRuntime interface { + InvokeAgent(ctx context.Context, in InvokeAgentInput) (*InvokeAgentResult, error) + Retrieve(ctx context.Context, in RetrieveInput) (*RetrieveResult, error) + RetrieveAndGenerate(ctx context.Context, in RetrieveAndGenerateInput) (*RetrieveAndGenerateResult, error) +} From 68a5b5cd79e8a47f4be1cbe2f75bc546dca4df13 Mon Sep 17 00:00:00 2001 From: Satyam Trivedi Date: Tue, 28 Jul 2026 18:00:52 +0530 Subject: [PATCH 2/5] =?UTF-8?q?fix(bedrock):=20address=20PR=20review=20?= =?UTF-8?q?=E2=80=94=20concurrency,=20immutability,=20routing,=20fidelity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Applies the review feedback on #214. Concurrency / immutability (data-race + aliasing): - Guardrail policies deep-copied on create/update and on version snapshot, so version snapshots are truly immutable and callers can't mutate stored state. - guardrailRecord guarded by a sync.RWMutex (draft/versions/nextVer); reads copy under the lock. go vet copylocks-clean. - Copy-on-write for in-place mutators that a new -race test surfaced: StopEvaluationJob, marketplace Update/Register, AR-policy Update. - copyMap on customization-job/custom-model read paths; copyBytes on Evaluation/Inference/PolicyDefinition []byte stores; copyRaw on bedrock-agent json.RawMessage configs. - New -race concurrency tests (guardrail + evaluation job). Routing: - bedrock-agent handler anchors /knowledgebases, /flows, /prompts via underPrefix so bucket paths like /flows-prod fall through to S3 (no shadowing); test added. Fidelity / consistency: - List endpoints (async invoke, import/copy/eval jobs, inference profiles, prompt routers, AR policies, marketplace) use SortedValues() for deterministic order. - Duplicate CreateMarketplaceModelEndpoint / CreatePromptRouter now return AlreadyExists instead of silently overwriting. - bedrock-agent-runtime decode uses errors.Is(io.EOF); dead *Preparing constants removed. Verified: go build/vet/gofmt clean, golangci-lint clean, go test ./... passing, go test -race on the bedrock packages clean, and an exhaustive 102-operation real-SDK end-to-end run against the full server passing. Co-Authored-By: Claude Opus 4.8 (1M context) --- providers/aws/bedrock/asyncinvoke.go | 2 +- providers/aws/bedrock/bedrock.go | 11 +- .../aws/bedrock/concurrency_race_test.go | 109 ++++++++++++++++++ providers/aws/bedrock/guardrail_copy.go | 64 ++++++++++ providers/aws/bedrock/guardrail_copy_test.go | 106 +++++++++++++++++ providers/aws/bedrock/guardrail_versions.go | 63 +++++++++- providers/aws/bedrock/jobs.go | 19 +-- providers/aws/bedrock/management.go | 42 ++++--- .../aws/bedrock/marketplace_agreements.go | 32 +++-- providers/aws/bedrock/registries.go | 26 +++-- providers/aws/bedrockagent/bedrockagent.go | 14 +++ providers/aws/bedrockagent/datasources.go | 4 +- providers/aws/bedrockagent/flows.go | 4 +- providers/aws/bedrockagent/knowledgebases.go | 8 +- providers/aws/bedrockagent/prompts.go | 4 +- server/aws/bedrockagent/handler.go | 23 ++-- server/aws/bedrockagent/sdk_roundtrip_test.go | 36 ++++++ server/aws/bedrockagentruntime/operations.go | 4 +- services/bedrockagent/driver/driver.go | 2 - 19 files changed, 500 insertions(+), 73 deletions(-) create mode 100644 providers/aws/bedrock/concurrency_race_test.go create mode 100644 providers/aws/bedrock/guardrail_copy.go create mode 100644 providers/aws/bedrock/guardrail_copy_test.go diff --git a/providers/aws/bedrock/asyncinvoke.go b/providers/aws/bedrock/asyncinvoke.go index 6193316c..574d3dfa 100644 --- a/providers/aws/bedrock/asyncinvoke.go +++ b/providers/aws/bedrock/asyncinvoke.go @@ -63,7 +63,7 @@ func (m *Mock) GetAsyncInvoke(_ context.Context, invocationARN string) (*driver. // ListAsyncInvokes lists all async invocations. func (m *Mock) ListAsyncInvokes(_ context.Context) ([]driver.AsyncInvoke, error) { - all := m.asyncInvokes.All() + all := m.asyncInvokes.SortedValues() out := make([]driver.AsyncInvoke, 0, len(all)) for _, inv := range all { diff --git a/providers/aws/bedrock/bedrock.go b/providers/aws/bedrock/bedrock.go index 7de3de33..2c09b62b 100644 --- a/providers/aws/bedrock/bedrock.go +++ b/providers/aws/bedrock/bedrock.go @@ -179,6 +179,7 @@ func (m *Mock) CreateModelCustomizationJob(_ context.Context, cfg driver.Customi func (m *Mock) GetModelCustomizationJob(_ context.Context, jobIdentifier string) (*driver.CustomizationJob, error) { if job, ok := m.jobs.Get(jobIdentifier); ok { result := *job + result.HyperParameters = copyMap(job.HyperParameters) return &result, nil } @@ -186,6 +187,7 @@ func (m *Mock) GetModelCustomizationJob(_ context.Context, jobIdentifier string) for _, job := range m.jobs.All() { if job.JobARN == jobIdentifier { result := *job + result.HyperParameters = copyMap(job.HyperParameters) return &result, nil } @@ -200,7 +202,9 @@ func (m *Mock) ListModelCustomizationJobs(_ context.Context) ([]driver.Customiza out := make([]driver.CustomizationJob, 0, len(all)) for _, job := range all { - out = append(out, *job) + result := *job + result.HyperParameters = copyMap(job.HyperParameters) + out = append(out, result) } return out, nil @@ -212,7 +216,9 @@ func (m *Mock) ListCustomModels(_ context.Context) ([]driver.CustomModel, error) out := make([]driver.CustomModel, 0, len(all)) for _, cm := range all { - out = append(out, *cm) + result := *cm + result.HyperParameters = copyMap(cm.HyperParameters) + out = append(out, result) } return out, nil @@ -226,6 +232,7 @@ func (m *Mock) GetCustomModel(_ context.Context, modelIdentifier string) (*drive } result := *cm + result.HyperParameters = copyMap(cm.HyperParameters) return &result, nil } diff --git a/providers/aws/bedrock/concurrency_race_test.go b/providers/aws/bedrock/concurrency_race_test.go new file mode 100644 index 00000000..32be3261 --- /dev/null +++ b/providers/aws/bedrock/concurrency_race_test.go @@ -0,0 +1,109 @@ +package bedrock + +import ( + "context" + "sync" + "testing" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// TestConcurrentGuardrailAccessRace hammers the guardrail record (versions slice +// + draft) with concurrent mutators and readers. Run under `go test -race` it +// guards against the data race fixed by the guardrailRecord mutex (M4) and the +// policy deep-copy (M3). +func TestConcurrentGuardrailAccessRace(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateGuardrail(ctx, bedrockdriver.GuardrailConfig{ + Name: "race-guard", + BlockedInputMessaging: "x", + BlockedOutputsMessaging: "y", + GuardrailPolicies: bedrockdriver.GuardrailPolicies{ + ContentPolicy: &bedrockdriver.GuardrailContentPolicy{ + Filters: []bedrockdriver.GuardrailContentFilter{{Type: "VIOLENCE", InputStrength: "HIGH", OutputStrength: "HIGH"}}, + }, + }, + }) + if err != nil { + t.Fatalf("CreateGuardrail: %v", err) + } + + const workers = 8 + + const iters = 50 + + var wg sync.WaitGroup + + for w := 0; w < workers; w++ { + wg.Add(1) + + go func(id int) { + defer wg.Done() + + for i := 0; i < iters; i++ { + switch id % 4 { + case 0: + _, _, _ = m.CreateGuardrailVersion(ctx, "race-guard", "v") + case 1: + _, _ = m.UpdateGuardrail(ctx, "race-guard", bedrockdriver.GuardrailConfig{ + Name: "race-guard", BlockedInputMessaging: "x2", BlockedOutputsMessaging: "y2", + }) + case 2: + _, _ = m.ListGuardrails(ctx, "race-guard") + default: + _, _ = m.GetGuardrail(ctx, "race-guard", "") + } + } + }(w) + } + + wg.Wait() +} + +// TestConcurrentEvalJobAccessRace hammers an evaluation job with concurrent Stop +// (mutating status through the stored pointer) and Get/List reads, catching the +// shared-pointer mutation class under `go test -race`. +func TestConcurrentEvalJobAccessRace(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateEvaluationJob(ctx, bedrockdriver.EvaluationJobConfig{ + JobName: "race-eval", + RoleARN: "arn:aws:iam::123456789012:role/bedrock", + EvaluationConfig: []byte(`{"automated":{}}`), + InferenceConfig: []byte(`{"models":[]}`), + OutputDataS3URI: "s3://bucket/eval/", + }) + if err != nil { + t.Fatalf("CreateEvaluationJob: %v", err) + } + + const workers = 8 + + const iters = 50 + + var wg sync.WaitGroup + + for w := 0; w < workers; w++ { + wg.Add(1) + + go func(id int) { + defer wg.Done() + + for i := 0; i < iters; i++ { + switch id % 3 { + case 0: + _ = m.StopEvaluationJob(ctx, "race-eval") + case 1: + _, _ = m.GetEvaluationJob(ctx, "race-eval") + default: + _, _ = m.ListEvaluationJobs(ctx) + } + } + }(w) + } + + wg.Wait() +} diff --git a/providers/aws/bedrock/guardrail_copy.go b/providers/aws/bedrock/guardrail_copy.go new file mode 100644 index 00000000..6312334a --- /dev/null +++ b/providers/aws/bedrock/guardrail_copy.go @@ -0,0 +1,64 @@ +package bedrock + +import ( + "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// deepCopyGuardrailPolicies returns a deep copy of the given guardrail +// policies, allocating fresh pointers for each non-nil sub-policy and cloning +// every slice. Nil sub-policies are left nil. This severs any aliasing between +// a caller's config, the stored draft, and immutable version snapshots. +func deepCopyGuardrailPolicies(p driver.GuardrailPolicies) driver.GuardrailPolicies { + var out driver.GuardrailPolicies + + if p.TopicPolicy != nil { + out.TopicPolicy = &driver.GuardrailTopicPolicy{ + Topics: copyTopics(p.TopicPolicy.Topics), + } + } + + if p.ContentPolicy != nil { + out.ContentPolicy = &driver.GuardrailContentPolicy{ + Filters: append([]driver.GuardrailContentFilter(nil), p.ContentPolicy.Filters...), + } + } + + if p.WordPolicy != nil { + out.WordPolicy = &driver.GuardrailWordPolicy{ + Words: append([]driver.GuardrailWord(nil), p.WordPolicy.Words...), + ManagedWordLists: append([]driver.GuardrailManagedWordList(nil), p.WordPolicy.ManagedWordLists...), + } + } + + if p.SensitiveInformationPolicy != nil { + out.SensitiveInformationPolicy = &driver.GuardrailSensitiveInformationPolicy{ + PiiEntities: append([]driver.GuardrailPiiEntity(nil), p.SensitiveInformationPolicy.PiiEntities...), + Regexes: append([]driver.GuardrailRegex(nil), p.SensitiveInformationPolicy.Regexes...), + } + } + + if p.ContextualGroundingPolicy != nil { + out.ContextualGroundingPolicy = &driver.GuardrailContextualGroundingPolicy{ + Filters: append([]driver.GuardrailContextualGroundingFilter(nil), p.ContextualGroundingPolicy.Filters...), + } + } + + return out +} + +// copyTopics clones a slice of topics, including each topic's Examples slice so +// the copy shares no backing array with the original. +func copyTopics(src []driver.GuardrailTopic) []driver.GuardrailTopic { + if src == nil { + return nil + } + + out := make([]driver.GuardrailTopic, len(src)) + + for i, t := range src { + t.Examples = append([]string(nil), t.Examples...) + out[i] = t + } + + return out +} diff --git a/providers/aws/bedrock/guardrail_copy_test.go b/providers/aws/bedrock/guardrail_copy_test.go new file mode 100644 index 00000000..5f4512d4 --- /dev/null +++ b/providers/aws/bedrock/guardrail_copy_test.go @@ -0,0 +1,106 @@ +package bedrock + +import ( + "context" + "testing" + + bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" +) + +// TestGuardrailPoliciesStoreDeepCopy verifies that CreateGuardrail deep-copies +// the caller's policy config: mutating the ORIGINAL cfg slices after create +// must not affect the stored draft. +func TestGuardrailPoliciesStoreDeepCopy(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + topics := []bedrockdriver.GuardrailTopic{ + {Name: "fiduciary", Definition: "financial advice", Examples: []string{"invest?"}, Type: "DENY"}, + } + filters := []bedrockdriver.GuardrailContentFilter{ + {Type: "HATE", InputStrength: "HIGH", OutputStrength: "MEDIUM"}, + } + + _, err := m.CreateGuardrail(ctx, bedrockdriver.GuardrailConfig{ + Name: "gr-immut", + BlockedInputMessaging: "in", + BlockedOutputsMessaging: "out", + GuardrailPolicies: bedrockdriver.GuardrailPolicies{ + TopicPolicy: &bedrockdriver.GuardrailTopicPolicy{Topics: topics}, + ContentPolicy: &bedrockdriver.GuardrailContentPolicy{Filters: filters}, + }, + }) + requireNoError(t, err) + + // Mutate the ORIGINAL caller-owned slices after the create call. + topics[0].Name = "MUTATED" + topics[0].Examples[0] = "MUTATED" + filters[0].InputStrength = "MUTATED" + + got, err := m.GetGuardrail(ctx, "gr-immut", "") + requireNoError(t, err) + + if got.TopicPolicy == nil || got.TopicPolicy.Topics[0].Name != "fiduciary" { + t.Fatalf("stored topic name aliased caller slice: %+v", got.TopicPolicy) + } + + if got.TopicPolicy.Topics[0].Examples[0] != "invest?" { + t.Fatalf("stored topic examples aliased caller slice: %+v", got.TopicPolicy.Topics[0].Examples) + } + + if got.ContentPolicy == nil || got.ContentPolicy.Filters[0].InputStrength != "HIGH" { + t.Fatalf("stored content filter aliased caller slice: %+v", got.ContentPolicy) + } +} + +// TestGuardrailVersionImmutableAgainstDraftEdits verifies that a numbered +// version snapshot is immutable against later edits to the DRAFT working copy. +func TestGuardrailVersionImmutableAgainstDraftEdits(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateGuardrail(ctx, bedrockdriver.GuardrailConfig{ + Name: "gr-ver-immut", + BlockedInputMessaging: "in", + BlockedOutputsMessaging: "out", + GuardrailPolicies: bedrockdriver.GuardrailPolicies{ + TopicPolicy: &bedrockdriver.GuardrailTopicPolicy{Topics: []bedrockdriver.GuardrailTopic{ + {Name: "orig", Definition: "d", Examples: []string{"e"}, Type: "DENY"}, + }}, + }, + }) + requireNoError(t, err) + + _, ver, err := m.CreateGuardrailVersion(ctx, "gr-ver-immut", "snapshot") + requireNoError(t, err) + assertEqual(t, "1", ver) + + // Update the DRAFT with entirely new policies. + _, err = m.UpdateGuardrail(ctx, "gr-ver-immut", bedrockdriver.GuardrailConfig{ + Name: "gr-ver-immut", + BlockedInputMessaging: "in", + BlockedOutputsMessaging: "out", + GuardrailPolicies: bedrockdriver.GuardrailPolicies{ + TopicPolicy: &bedrockdriver.GuardrailTopicPolicy{Topics: []bedrockdriver.GuardrailTopic{ + {Name: "changed", Definition: "d2", Examples: []string{"e2"}, Type: "DENY"}, + }}, + }, + }) + requireNoError(t, err) + + // The version snapshot must still carry the original policy values. + snap, err := m.GetGuardrail(ctx, "gr-ver-immut", ver) + requireNoError(t, err) + + if snap.TopicPolicy == nil || snap.TopicPolicy.Topics[0].Name != "orig" { + t.Fatalf("version snapshot mutated by draft edit: %+v", snap.TopicPolicy) + } + + // And the DRAFT reflects the update, confirming they are distinct graphs. + draft, err := m.GetGuardrail(ctx, "gr-ver-immut", "DRAFT") + requireNoError(t, err) + + if draft.TopicPolicy.Topics[0].Name != "changed" { + t.Fatalf("draft did not pick up the update: %+v", draft.TopicPolicy) + } +} diff --git a/providers/aws/bedrock/guardrail_versions.go b/providers/aws/bedrock/guardrail_versions.go index a0a9ad44..6928d828 100644 --- a/providers/aws/bedrock/guardrail_versions.go +++ b/providers/aws/bedrock/guardrail_versions.go @@ -3,6 +3,7 @@ package bedrock import ( "context" "strconv" + "sync" "github.com/stackshy/cloudemu/v2/errors" "github.com/stackshy/cloudemu/v2/services/bedrock/driver" @@ -11,13 +12,20 @@ import ( // guardrailRecord holds a guardrail's mutable working copy (the "DRAFT" // version) alongside its immutable, numbered version snapshots. Records are // stored in m.guardrails keyed by the guardrail name. +// +// mu guards draft, versions and nextVer. memstore only serializes access to +// its map, not to the record it points at, so every reader and mutator of a +// *guardrailRecord must hold mu. The record therefore contains a sync.RWMutex +// and MUST NOT be copied by value — always pass and store it as a pointer. type guardrailRecord struct { + mu sync.RWMutex draft *driver.Guardrail versions []*driver.Guardrail // numbered snapshots ("1", "2", ...) in creation order nextVer int // monotonic next version number; never reused after delete } -// version returns the snapshot for the given numbered version, or nil. +// version returns the snapshot for the given numbered version, or nil. Callers +// must hold rec.mu (read or write); it performs no locking of its own. func (rec *guardrailRecord) version(v string) *driver.Guardrail { for _, g := range rec.versions { if g.Version == v { @@ -28,6 +36,52 @@ func (rec *guardrailRecord) version(v string) *driver.Guardrail { return nil } +// draftSnapshot returns a deep, aliasing-free copy of the DRAFT working copy. +func (rec *guardrailRecord) draftSnapshot() driver.Guardrail { + rec.mu.RLock() + defer rec.mu.RUnlock() + + return cloneGuardrail(rec.draft) +} + +// versionSnapshot returns a deep copy of the given numbered version and true, +// or the zero value and false if no such version exists. +func (rec *guardrailRecord) versionSnapshot(v string) (driver.Guardrail, bool) { + rec.mu.RLock() + defer rec.mu.RUnlock() + + g := rec.version(v) + if g == nil { + return driver.Guardrail{}, false + } + + return cloneGuardrail(g), true +} + +// allSnapshots returns deep copies of the DRAFT plus every numbered version. +func (rec *guardrailRecord) allSnapshots() []driver.Guardrail { + rec.mu.RLock() + defer rec.mu.RUnlock() + + out := make([]driver.Guardrail, 0, len(rec.versions)+1) + out = append(out, cloneGuardrail(rec.draft)) + + for _, v := range rec.versions { + out = append(out, cloneGuardrail(v)) + } + + return out +} + +// cloneGuardrail returns a value copy of g with its policy object graph +// deep-copied, so the result shares no mutable state with the stored record. +func cloneGuardrail(g *driver.Guardrail) driver.Guardrail { + result := *g + result.GuardrailPolicies = deepCopyGuardrailPolicies(g.GuardrailPolicies) + + return result +} + // CreateGuardrailVersion snapshots the current DRAFT into a new immutable, // numbered version and returns the guardrail ID and the assigned version. func (m *Mock) CreateGuardrailVersion( @@ -38,10 +92,15 @@ func (m *Mock) CreateGuardrailVersion( return "", "", errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } + rec.mu.Lock() + defer rec.mu.Unlock() + ver := strconv.Itoa(rec.nextVer) rec.nextVer++ - snapshot := *rec.draft + // Deep-copy the draft so the snapshot is immutable against later draft + // edits (UpdateGuardrail mutates the same *draft in place). + snapshot := cloneGuardrail(rec.draft) snapshot.Version = ver if description != "" { diff --git a/providers/aws/bedrock/jobs.go b/providers/aws/bedrock/jobs.go index 46414deb..7fec38b9 100644 --- a/providers/aws/bedrock/jobs.go +++ b/providers/aws/bedrock/jobs.go @@ -78,7 +78,7 @@ func (m *Mock) GetModelImportJob(_ context.Context, jobIdentifier string) (*driv // ListModelImportJobs lists all import jobs. func (m *Mock) ListModelImportJobs(_ context.Context) ([]driver.ModelImportJob, error) { - all := m.importJobs.All() + all := m.importJobs.SortedValues() out := make([]driver.ModelImportJob, 0, len(all)) for _, job := range all { @@ -136,7 +136,7 @@ func (m *Mock) GetModelCopyJob(_ context.Context, jobARN string) (*driver.ModelC // ListModelCopyJobs lists all copy jobs. func (m *Mock) ListModelCopyJobs(_ context.Context) ([]driver.ModelCopyJob, error) { - all := m.copyJobs.All() + all := m.copyJobs.SortedValues() out := make([]driver.ModelCopyJob, 0, len(all)) for _, job := range all { @@ -179,8 +179,8 @@ func (m *Mock) CreateEvaluationJob(_ context.Context, cfg driver.EvaluationJobCo JobType: evaluationJobType(cfg.EvaluationConfig), ApplicationType: cfg.ApplicationType, RoleARN: cfg.RoleARN, - EvaluationConfig: cfg.EvaluationConfig, - InferenceConfig: cfg.InferenceConfig, + EvaluationConfig: copyBytes(cfg.EvaluationConfig), + InferenceConfig: copyBytes(cfg.InferenceConfig), OutputDataS3URI: cfg.OutputDataS3URI, JobDescription: cfg.JobDescription, CustomerEncryptionKeyID: cfg.CustomerEncryptionKeyID, @@ -210,7 +210,7 @@ func (m *Mock) GetEvaluationJob(_ context.Context, jobIdentifier string) (*drive // ListEvaluationJobs lists all evaluation jobs. func (m *Mock) ListEvaluationJobs(_ context.Context) ([]driver.EvaluationJob, error) { - all := m.evalJobs.All() + all := m.evalJobs.SortedValues() out := make([]driver.EvaluationJob, 0, len(all)) for _, job := range all { @@ -227,9 +227,12 @@ func (m *Mock) StopEvaluationJob(_ context.Context, jobIdentifier string) error return errors.Newf(errors.NotFound, "evaluation job %q not found", jobIdentifier) } - job.Status = driver.JobStopped - job.LastModifiedTime = m.now() - m.evalJobs.Set(job.JobName, job) + // Copy-on-write: never mutate the stored pointer in place, so concurrent + // Get/List readers (which copy the stored value) can't race the write. + updated := *job + updated.Status = driver.JobStopped + updated.LastModifiedTime = m.now() + m.evalJobs.Set(updated.JobName, &updated) return nil } diff --git a/providers/aws/bedrock/management.go b/providers/aws/bedrock/management.go index 8309b4f0..fcee0019 100644 --- a/providers/aws/bedrock/management.go +++ b/providers/aws/bedrock/management.go @@ -39,7 +39,7 @@ func (m *Mock) CreateGuardrail(_ context.Context, cfg driver.GuardrailConfig) (* BlockedOutputsMessaging: cfg.BlockedOutputsMessaging, CreatedAt: now, UpdatedAt: now, - GuardrailPolicies: cfg.GuardrailPolicies, + GuardrailPolicies: deepCopyGuardrailPolicies(cfg.GuardrailPolicies), } if cfg.KMSKeyID != "" { @@ -62,14 +62,16 @@ func (m *Mock) GetGuardrail(_ context.Context, identifier, version string) (*dri return nil, errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } - g := rec.draft if version != "" && version != guardrailDraftVersion { - if g = rec.version(version); g == nil { + result, ok := rec.versionSnapshot(version) + if !ok { return nil, errors.Newf(errors.NotFound, "guardrail %q version %q not found", identifier, version) } + + return &result, nil } - result := *g + result := rec.draftSnapshot() return &result, nil } @@ -84,21 +86,14 @@ func (m *Mock) ListGuardrails(_ context.Context, identifier string) ([]driver.Gu return nil, errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } - out := make([]driver.Guardrail, 0, len(rec.versions)+1) - out = append(out, *rec.draft) - - for _, v := range rec.versions { - out = append(out, *v) - } - - return out, nil + return rec.allSnapshots(), nil } all := m.guardrails.All() out := make([]driver.Guardrail, 0, len(all)) for _, rec := range all { - out = append(out, *rec.draft) + out = append(out, rec.draftSnapshot()) } return out, nil @@ -114,24 +109,26 @@ func (m *Mock) UpdateGuardrail(_ context.Context, identifier string, cfg driver. return nil, errors.Newf(errors.NotFound, "guardrail %q not found", identifier) } + rec.mu.Lock() g := rec.draft oldName := g.Name g.Name = orDefault(cfg.Name, g.Name) g.Description = cfg.Description g.BlockedInputMessaging = orDefault(cfg.BlockedInputMessaging, g.BlockedInputMessaging) g.BlockedOutputsMessaging = orDefault(cfg.BlockedOutputsMessaging, g.BlockedOutputsMessaging) - g.GuardrailPolicies = cfg.GuardrailPolicies + g.GuardrailPolicies = deepCopyGuardrailPolicies(cfg.GuardrailPolicies) g.UpdatedAt = m.now() + newName := g.Name + result := cloneGuardrail(g) + rec.mu.Unlock() // Records are keyed by name; re-key when an update renames one so lookups // by the new name keep working. - if g.Name != oldName { + if newName != oldName { m.guardrails.Delete(oldName) - m.guardrails.Set(g.Name, rec) + m.guardrails.Set(newName, rec) } - result := *g - return &result, nil } @@ -145,11 +142,18 @@ func (m *Mock) DeleteGuardrail(_ context.Context, identifier, version string) er } if version == "" { - m.guardrails.Delete(rec.draft.Name) + rec.mu.RLock() + name := rec.draft.Name + rec.mu.RUnlock() + + m.guardrails.Delete(name) return nil } + rec.mu.Lock() + defer rec.mu.Unlock() + for i, v := range rec.versions { if v.Version == version { rec.versions = append(rec.versions[:i], rec.versions[i+1:]...) diff --git a/providers/aws/bedrock/marketplace_agreements.go b/providers/aws/bedrock/marketplace_agreements.go index c280da76..0e3299ad 100644 --- a/providers/aws/bedrock/marketplace_agreements.go +++ b/providers/aws/bedrock/marketplace_agreements.go @@ -30,6 +30,10 @@ func (m *Mock) CreateMarketplaceModelEndpoint( now := m.now() arn := idgen.AWSARN("sagemaker", m.opts.Region, m.opts.AccountID, "endpoint/"+cfg.EndpointName) + if m.marketplaceEndpoints.Has(arn) { + return nil, errors.Newf(errors.AlreadyExists, "marketplace model endpoint %q already exists", cfg.EndpointName) + } + endpoint := &driver.MarketplaceEndpoint{ EndpointARN: arn, ModelSourceIdentifier: cfg.ModelSourceIdentifier, @@ -61,7 +65,7 @@ func (m *Mock) GetMarketplaceModelEndpoint(_ context.Context, endpointARN string // ListMarketplaceModelEndpoints lists all marketplace model endpoints. func (m *Mock) ListMarketplaceModelEndpoints(_ context.Context) ([]driver.MarketplaceEndpoint, error) { - all := m.marketplaceEndpoints.All() + all := m.marketplaceEndpoints.SortedValues() out := make([]driver.MarketplaceEndpoint, 0, len(all)) for _, endpoint := range all { @@ -75,7 +79,7 @@ func (m *Mock) ListMarketplaceModelEndpoints(_ context.Context) ([]driver.Market func (m *Mock) UpdateMarketplaceModelEndpoint( _ context.Context, endpointARN string, endpointConfig []byte, ) (*driver.MarketplaceEndpoint, error) { - endpoint, ok := m.marketplaceEndpoints.Get(endpointARN) + stored, ok := m.marketplaceEndpoints.Get(endpointARN) if !ok { return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) } @@ -84,11 +88,13 @@ func (m *Mock) UpdateMarketplaceModelEndpoint( return nil, errors.New(errors.InvalidArgument, "endpointConfig is required") } - endpoint.EndpointConfig = copyBytes(endpointConfig) - endpoint.UpdatedAt = m.now() - m.marketplaceEndpoints.Set(endpointARN, endpoint) + // Copy-on-write: mutate a copy so concurrent readers never race the write. + updated := *stored + updated.EndpointConfig = copyBytes(endpointConfig) + updated.UpdatedAt = m.now() + m.marketplaceEndpoints.Set(endpointARN, &updated) - result := *endpoint + result := updated return &result, nil } @@ -113,17 +119,19 @@ func (m *Mock) RegisterMarketplaceModelEndpoint( return nil, errors.New(errors.InvalidArgument, "modelSourceIdentifier is required") } - endpoint, ok := m.marketplaceEndpoints.Get(endpointIdentifier) + stored, ok := m.marketplaceEndpoints.Get(endpointIdentifier) if !ok { return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointIdentifier) } - endpoint.ModelSourceIdentifier = modelSourceIdentifier - endpoint.Status = driver.MarketplaceEndpointStatusRegistered - endpoint.UpdatedAt = m.now() - m.marketplaceEndpoints.Set(endpointIdentifier, endpoint) + // Copy-on-write: mutate a copy so concurrent readers never race the write. + updated := *stored + updated.ModelSourceIdentifier = modelSourceIdentifier + updated.Status = driver.MarketplaceEndpointStatusRegistered + updated.UpdatedAt = m.now() + m.marketplaceEndpoints.Set(endpointIdentifier, &updated) - result := *endpoint + result := updated return &result, nil } diff --git a/providers/aws/bedrock/registries.go b/providers/aws/bedrock/registries.go index 3d4078f3..c546fdf0 100644 --- a/providers/aws/bedrock/registries.go +++ b/providers/aws/bedrock/registries.go @@ -66,7 +66,7 @@ func (m *Mock) GetInferenceProfile(_ context.Context, identifier string) (*drive // ListInferenceProfiles lists all inference profiles. func (m *Mock) ListInferenceProfiles(_ context.Context) ([]driver.InferenceProfile, error) { - all := m.inferenceProfiles.All() + all := m.inferenceProfiles.SortedValues() out := make([]driver.InferenceProfile, 0, len(all)) for _, p := range all { @@ -113,6 +113,12 @@ func (m *Mock) CreatePromptRouter(_ context.Context, cfg driver.PromptRouterConf return nil, errors.New(errors.InvalidArgument, "fallbackModel.modelArn is required") } + for _, existing := range m.promptRouters.SortedValues() { + if existing.Name == cfg.Name { + return nil, errors.Newf(errors.AlreadyExists, "prompt router %q already exists", cfg.Name) + } + } + now := m.now() id := idgen.GenerateID("") arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "prompt-router/"+id) @@ -151,7 +157,7 @@ func (m *Mock) GetPromptRouter(_ context.Context, promptRouterARN string) (*driv // ListPromptRouters lists all prompt routers. func (m *Mock) ListPromptRouters(_ context.Context) ([]driver.PromptRouter, error) { - all := m.promptRouters.All() + all := m.promptRouters.SortedValues() out := make([]driver.PromptRouter, 0, len(all)) for _, router := range all { @@ -196,7 +202,7 @@ func (m *Mock) CreateAutomatedReasoningPolicy( Version: driver.AutomatedReasoningPolicyVersionDraft, DefinitionHash: definitionHash(cfg.PolicyDefinition), Description: cfg.Description, - PolicyDefinition: cfg.PolicyDefinition, + PolicyDefinition: copyBytes(cfg.PolicyDefinition), CreatedAt: now, UpdatedAt: now, } @@ -226,7 +232,7 @@ func (m *Mock) GetAutomatedReasoningPolicy(_ context.Context, policyARN string) // ListAutomatedReasoningPolicies lists all automated reasoning policies. func (m *Mock) ListAutomatedReasoningPolicies(_ context.Context) ([]driver.AutomatedReasoningPolicy, error) { - all := m.arPolicies.All() + all := m.arPolicies.SortedValues() out := make([]driver.AutomatedReasoningPolicy, 0, len(all)) for _, policy := range all { @@ -241,11 +247,15 @@ func (m *Mock) ListAutomatedReasoningPolicies(_ context.Context) ([]driver.Autom func (m *Mock) UpdateAutomatedReasoningPolicy( _ context.Context, policyARN string, upd driver.AutomatedReasoningPolicyUpdate, ) (*driver.AutomatedReasoningPolicy, error) { - policy, ok := m.arPolicies.Get(policyARN) + stored, ok := m.arPolicies.Get(policyARN) if !ok { return nil, errors.Newf(errors.NotFound, "automated reasoning policy %q not found", policyARN) } + // Copy-on-write: mutate a copy, not the stored pointer, so concurrent + // Get/List readers never race the write. + policy := *stored + if upd.Name != "" { policy.Name = upd.Name } @@ -255,14 +265,14 @@ func (m *Mock) UpdateAutomatedReasoningPolicy( } if len(upd.PolicyDefinition) != 0 { - policy.PolicyDefinition = upd.PolicyDefinition + policy.PolicyDefinition = copyBytes(upd.PolicyDefinition) } policy.DefinitionHash = definitionHash(policy.PolicyDefinition) policy.UpdatedAt = m.now() - m.arPolicies.Set(policyARN, policy) + m.arPolicies.Set(policyARN, &policy) - result := *policy + result := policy return &result, nil } diff --git a/providers/aws/bedrockagent/bedrockagent.go b/providers/aws/bedrockagent/bedrockagent.go index df3439e7..262a3831 100644 --- a/providers/aws/bedrockagent/bedrockagent.go +++ b/providers/aws/bedrockagent/bedrockagent.go @@ -5,6 +5,7 @@ package bedrockagent import ( + "encoding/json" "time" "github.com/stackshy/cloudemu/v2/config" @@ -51,3 +52,16 @@ func New(opts *config.Options) *Mock { func (m *Mock) now() string { return m.opts.Clock.Now().UTC().Format(time.RFC3339) } + +// copyRaw returns a defensive copy of a json.RawMessage so stored resources do +// not alias caller-owned buffers. nil maps to nil. +func copyRaw(b json.RawMessage) json.RawMessage { + if b == nil { + return nil + } + + out := make(json.RawMessage, len(b)) + copy(out, b) + + return out +} diff --git a/providers/aws/bedrockagent/datasources.go b/providers/aws/bedrockagent/datasources.go index cd83f0bf..9faa05a0 100644 --- a/providers/aws/bedrockagent/datasources.go +++ b/providers/aws/bedrockagent/datasources.go @@ -33,7 +33,7 @@ func (m *Mock) CreateDataSource(_ context.Context, cfg driver.DataSourceConfig) Description: cfg.Description, Status: driver.DataSourceAvailable, DataDeletionPolicy: cfg.DataDeletionPolicy, - DataSourceConfiguration: cfg.DataSourceConfiguration, + DataSourceConfiguration: copyRaw(cfg.DataSourceConfiguration), CreatedAt: now, UpdatedAt: now, } @@ -86,7 +86,7 @@ func (m *Mock) UpdateDataSource(_ context.Context, cfg driver.DataSourceConfig, updated.UpdatedAt = m.now() if len(cfg.DataSourceConfiguration) != 0 { - updated.DataSourceConfiguration = cfg.DataSourceConfiguration + updated.DataSourceConfiguration = copyRaw(cfg.DataSourceConfiguration) } m.dataSource.Set(dsID, &updated) diff --git a/providers/aws/bedrockagent/flows.go b/providers/aws/bedrockagent/flows.go index 37a6f5d0..170a94b9 100644 --- a/providers/aws/bedrockagent/flows.go +++ b/providers/aws/bedrockagent/flows.go @@ -30,7 +30,7 @@ func (m *Mock) CreateFlow(_ context.Context, cfg driver.FlowConfig) (*driver.Flo Status: driver.FlowNotPrepared, Version: driver.DraftVersion, CustomerEncryptionKeyArn: cfg.CustomerEncryptionKeyArn, - Definition: cfg.Definition, + Definition: copyRaw(cfg.Definition), CreatedAt: now, UpdatedAt: now, } @@ -82,7 +82,7 @@ func (m *Mock) UpdateFlow(_ context.Context, id string, cfg driver.FlowConfig) ( updated.UpdatedAt = m.now() if len(cfg.Definition) != 0 { - updated.Definition = cfg.Definition + updated.Definition = copyRaw(cfg.Definition) } m.flows.Set(id, &updated) diff --git a/providers/aws/bedrockagent/knowledgebases.go b/providers/aws/bedrockagent/knowledgebases.go index 7f085994..9b68d59f 100644 --- a/providers/aws/bedrockagent/knowledgebases.go +++ b/providers/aws/bedrockagent/knowledgebases.go @@ -30,8 +30,8 @@ func (m *Mock) CreateKnowledgeBase(_ context.Context, cfg driver.KnowledgeBaseCo RoleArn: cfg.RoleArn, Description: cfg.Description, Status: driver.KnowledgeBaseActive, - KnowledgeBaseConfiguration: cfg.KnowledgeBaseConfiguration, - StorageConfiguration: cfg.StorageConfiguration, + KnowledgeBaseConfiguration: copyRaw(cfg.KnowledgeBaseConfiguration), + StorageConfiguration: copyRaw(cfg.StorageConfiguration), CreatedAt: now, UpdatedAt: now, } @@ -82,11 +82,11 @@ func (m *Mock) UpdateKnowledgeBase(_ context.Context, id string, cfg driver.Know updated.UpdatedAt = m.now() if len(cfg.KnowledgeBaseConfiguration) != 0 { - updated.KnowledgeBaseConfiguration = cfg.KnowledgeBaseConfiguration + updated.KnowledgeBaseConfiguration = copyRaw(cfg.KnowledgeBaseConfiguration) } if len(cfg.StorageConfiguration) != 0 { - updated.StorageConfiguration = cfg.StorageConfiguration + updated.StorageConfiguration = copyRaw(cfg.StorageConfiguration) } m.knowledge.Set(id, &updated) diff --git a/providers/aws/bedrockagent/prompts.go b/providers/aws/bedrockagent/prompts.go index 5a29e07f..4f27e0d0 100644 --- a/providers/aws/bedrockagent/prompts.go +++ b/providers/aws/bedrockagent/prompts.go @@ -26,7 +26,7 @@ func (m *Mock) CreatePrompt(_ context.Context, cfg driver.PromptConfig) (*driver Version: driver.DraftVersion, DefaultVariant: cfg.DefaultVariant, CustomerEncryptionKeyArn: cfg.CustomerEncryptionKeyArn, - Variants: cfg.Variants, + Variants: copyRaw(cfg.Variants), CreatedAt: now, UpdatedAt: now, } @@ -77,7 +77,7 @@ func (m *Mock) UpdatePrompt(_ context.Context, id string, cfg driver.PromptConfi updated.UpdatedAt = m.now() if len(cfg.Variants) != 0 { - updated.Variants = cfg.Variants + updated.Variants = copyRaw(cfg.Variants) } m.prompts.Set(id, &updated) diff --git a/server/aws/bedrockagent/handler.go b/server/aws/bedrockagent/handler.go index 2c769212..53d29adf 100644 --- a/server/aws/bedrockagent/handler.go +++ b/server/aws/bedrockagent/handler.go @@ -81,10 +81,17 @@ func New(drv badriver.BedrockAgent) *Handler { func (*Handler) Matches(r *http.Request) bool { p := r.URL.Path - return strings.HasPrefix(p, prefixAgents+"/") || - strings.HasPrefix(p, prefixKB) || - strings.HasPrefix(p, prefixFlows) || - strings.HasPrefix(p, prefixPrompts) + return underPrefix(p, prefixAgents) || + underPrefix(p, prefixKB) || + underPrefix(p, prefixFlows) || + underPrefix(p, prefixPrompts) +} + +// underPrefix reports whether p equals pre or is a child path of pre. It +// anchors bare prefixes so bucket-style paths (e.g. "/flows-prod") fall +// through to later handlers instead of being swallowed here. +func underPrefix(p, pre string) bool { + return p == pre || strings.HasPrefix(p, pre+"/") } // ServeHTTP routes by URL prefix. @@ -92,13 +99,13 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { p := r.URL.Path switch { - case strings.HasPrefix(p, prefixAgents+"/"): + case underPrefix(p, prefixAgents): h.serveAgents(w, r, segments(p, prefixAgents)) - case strings.HasPrefix(p, prefixKB): + case underPrefix(p, prefixKB): h.serveKnowledgeBases(w, r, segments(p, prefixKB)) - case strings.HasPrefix(p, prefixFlows): + case underPrefix(p, prefixFlows): h.serveFlows(w, r, segments(p, prefixFlows)) - case strings.HasPrefix(p, prefixPrompts): + case underPrefix(p, prefixPrompts): h.servePrompts(w, r, segments(p, prefixPrompts)) default: notFound(w, p) diff --git a/server/aws/bedrockagent/sdk_roundtrip_test.go b/server/aws/bedrockagent/sdk_roundtrip_test.go index e9a59d9b..82a3d2b4 100644 --- a/server/aws/bedrockagent/sdk_roundtrip_test.go +++ b/server/aws/bedrockagent/sdk_roundtrip_test.go @@ -3,6 +3,7 @@ package bedrockagent_test import ( "context" "errors" + "net/http" "net/http/httptest" "testing" @@ -344,6 +345,41 @@ func TestSDKPromptLifecycle(t *testing.T) { assertNotFound(t, err) } +// TestMatchesAnchorsPrefixes guards the M2 fix: bucket-style paths that merely +// share a prefix (e.g. "/flows-prod") must NOT be claimed by the Bedrock Agent +// handler, so they fall through to the S3 catch-all, while the documented +// collection/item shapes still match. +func TestMatchesAnchorsPrefixes(t *testing.T) { + h := serverba.New(providerba.New(config.NewOptions())) + + cases := []struct { + path string + want bool + }{ + // Bucket-style paths must fall through to S3. + {"/flows-prod", false}, + {"/promptsdb", false}, + {"/knowledgebases-archive", false}, + {"/agents-backup", false}, + // Documented collection and item shapes must still be claimed. + {"/agents/", true}, + {"/knowledgebases", true}, + {"/knowledgebases/kb-123", true}, + {"/knowledgebases/kb-123/datasources/ds-1", true}, + {"/flows/", true}, + {"/flows/flow-1/", true}, + {"/prompts/", true}, + {"/prompts/prompt-1/", true}, + } + + for _, tc := range cases { + req := httptest.NewRequest(http.MethodGet, tc.path, nil) + if got := h.Matches(req); got != tc.want { + t.Errorf("Matches(%q) = %v, want %v", tc.path, got, tc.want) + } + } +} + func assertNotFound(t *testing.T, err error) { t.Helper() diff --git a/server/aws/bedrockagentruntime/operations.go b/server/aws/bedrockagentruntime/operations.go index c9bac610..d342e8b5 100644 --- a/server/aws/bedrockagentruntime/operations.go +++ b/server/aws/bedrockagentruntime/operations.go @@ -2,6 +2,8 @@ package bedrockagentruntime import ( "encoding/json" + "errors" + "io" "net/http" bedrockagentruntimedriver "github.com/stackshy/cloudemu/v2/services/bedrockagentruntime/driver" @@ -117,7 +119,7 @@ func decodeJSON(w http.ResponseWriter, r *http.Request, v any) bool { func decodeJSONAllowEmpty(w http.ResponseWriter, r *http.Request, v any) bool { r.Body = http.MaxBytesReader(w, r.Body, maxBodyBytes) - if err := json.NewDecoder(r.Body).Decode(v); err != nil && err.Error() != "EOF" { + if err := json.NewDecoder(r.Body).Decode(v); err != nil && !errors.Is(err, io.EOF) { writeError(w, http.StatusBadRequest, "ValidationException", "invalid JSON: "+err.Error()) return false diff --git a/services/bedrockagent/driver/driver.go b/services/bedrockagent/driver/driver.go index def26de8..b98b8c42 100644 --- a/services/bedrockagent/driver/driver.go +++ b/services/bedrockagent/driver/driver.go @@ -11,7 +11,6 @@ import ( // Agent lifecycle status values (UPPER_SNAKE, per AgentStatus). const ( AgentNotPrepared = "NOT_PREPARED" - AgentPreparing = "PREPARING" AgentPrepared = "PREPARED" ) @@ -38,7 +37,6 @@ const ( // Flow status values (PascalCase, per FlowStatus). const ( FlowNotPrepared = "NotPrepared" - FlowPreparing = "Preparing" FlowPrepared = "Prepared" ) From ce487962321bcc34f2182c465abea75bff86c8ad Mon Sep 17 00:00:00 2001 From: Satyam Trivedi Date: Wed, 29 Jul 2026 11:34:37 +0530 Subject: [PATCH 3/5] fix(bedrock): address re-review fidelity items + arch pattern gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the second/third PR reviews on #214. Fidelity: - Evaluation jobs are now created InProgress (evaluation is long-running with no synchronous artifact, unlike import/copy), so StopEvaluationJob is meaningful; Stop rejects a non-InProgress job with FailedPrecondition, which the bedrock error mapper now surfaces as ConflictException (409), matching real AWS. - DeregisterMarketplaceModelEndpoint removes the Bedrock registration record, so a subsequent Get returns NotFound (the underlying SageMaker endpoint is unmodeled) — matching AWS instead of a success no-op. - ApplyGuardrail validates a requested numbered GuardrailVersion (NotFound if the version doesn't exist); "" / "DRAFT" resolve to the working draft. Aliasing/consistency: - GetInferenceProfile / GetPromptRouter / GetAutomatedReasoningPolicy (and their List) copy the Models slice / PolicyDefinition bytes out on read via clone helpers, so callers can't mutate stored state through a returned value. Pattern conformance: - Added the missing services/bedrockagent portable-layer test. - docs/services.md now lists bedrock-agent (+ bedrock-agent-runtime) and the operation-count table reflects the full surface. Tests updated for the corrected eval-job lifecycle and marketplace deregister semantics. Verified: go build/vet/gofmt/golangci-lint clean, go test ./... passing (174 pkgs), go test -race on the bedrock packages clean, and an exhaustive 107-operation real-SDK end-to-end run against the full server passing (including the changed behaviors). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/services.md | 7 +- .../aws/bedrock/asyncinvoke_jobs_test.go | 5 +- .../aws/bedrock/counttokens_applyguardrail.go | 11 ++- providers/aws/bedrock/jobs.go | 16 ++- .../aws/bedrock/marketplace_agreements.go | 9 +- .../bedrock/marketplace_agreements_test.go | 15 ++- providers/aws/bedrock/registries.go | 41 ++++++-- server/aws/bedrock/errors.go | 4 +- .../bedrock/sdk_roundtrip_asyncjobs_test.go | 4 +- .../bedrock/sdk_roundtrip_marketplace_test.go | 26 ++++- services/bedrockagent/bedrockagent_test.go | 99 +++++++++++++++++++ 11 files changed, 209 insertions(+), 28 deletions(-) create mode 100644 services/bedrockagent/bedrockagent_test.go diff --git a/docs/services.md b/docs/services.md index 4c63fa93..4f92cdec 100644 --- a/docs/services.md +++ b/docs/services.md @@ -25,7 +25,7 @@ This document lists every service and operation available in CloudEmu across all | 17 | Relational Database | `rds` (+ Aurora/Neptune/DocumentDB engines), `redshift` | `azuresql`, `postgresflex`, `mysqlflex` | `cloudsql` | | 18 | Kubernetes | `eks` + shared `services/kubernetes/` | `aks` + shared `services/kubernetes/` | `gke` + shared `services/kubernetes/` | | 19 | Resource Discovery | `resourceexplorer2` + `resourcegroupstaggingapi` | `resourcegraph` | `cloudasset` | -| 20 | Generative AI | `bedrock` (+ `bedrock-runtime`) | — | — | +| 20 | Generative AI | `bedrock` (+ `bedrock-runtime`), `bedrock-agent` (+ `bedrock-agent-runtime`) | — | — | | 21 | Databricks | — | `databricks` | — | | 22 | Machine Learning | `sagemaker` (+ `sagemaker-runtime`) | `azureai` (CognitiveServices + MachineLearningServices) | `vertexai` | | 23 | AI Search | — | `azuresearch` (Microsoft.Search) | — | @@ -1638,13 +1638,14 @@ still sees success. | Kubernetes — GCP GKE (control plane) | 26 | | Kubernetes — data plane (8 resources × 7 verbs incl. Watch) | 56 | | Resource Discovery (engine + AWS + Azure + GCP handlers) | 26 | -| Generative AI — AWS Bedrock | 22 | +| Generative AI — AWS Bedrock (control plane + runtime) | 65 | +| Generative AI — AWS Bedrock Agent (control plane + runtime) | 32 | | Databricks — Azure (control + data plane) | 52 | | Machine Learning — AWS SageMaker (control plane + runtime) | 121 | | Machine Learning — Azure AI (CognitiveServices + MachineLearningServices + data plane) | 92 | | Machine Learning — GCP Vertex AI (Go API/driver) | 128 | | AI Search — Azure AI Search (control + data plane) | 53 | -| **Grand Total** | **972** (+12 optional) | +| **Grand Total** | **1047** (+12 optional) | Optional operations are capabilities a driver may implement but is not required to; see the sections marked "optional capability". They are counted separately diff --git a/providers/aws/bedrock/asyncinvoke_jobs_test.go b/providers/aws/bedrock/asyncinvoke_jobs_test.go index 80fd1bc3..46af50b4 100644 --- a/providers/aws/bedrock/asyncinvoke_jobs_test.go +++ b/providers/aws/bedrock/asyncinvoke_jobs_test.go @@ -139,7 +139,7 @@ func TestEvaluationJobLifecycle(t *testing.T) { }) requireNoError(t, err) assertNotEmpty(t, job.JobARN) - assertEqual(t, bedrockdriver.JobCompleted, job.Status) + assertEqual(t, bedrockdriver.JobInProgress, job.Status) assertEqual(t, bedrockdriver.EvaluationTypeAutomated, job.JobType) got, err := m.GetEvaluationJob(ctx, "eval-1") @@ -160,6 +160,9 @@ func TestEvaluationJobLifecycle(t *testing.T) { requireNoError(t, err) assertEqual(t, bedrockdriver.JobStopped, stopped.Status) + // Stopping a job that is no longer in progress is rejected. + assertError(t, m.StopEvaluationJob(ctx, "eval-1"), true) + assertError(t, m.StopEvaluationJob(ctx, "missing"), true) } diff --git a/providers/aws/bedrock/counttokens_applyguardrail.go b/providers/aws/bedrock/counttokens_applyguardrail.go index 6fd2c3b0..59e37ade 100644 --- a/providers/aws/bedrock/counttokens_applyguardrail.go +++ b/providers/aws/bedrock/counttokens_applyguardrail.go @@ -37,10 +37,19 @@ func (m *Mock) ApplyGuardrail(_ context.Context, in driver.ApplyGuardrailInput) return nil, errors.New(errors.InvalidArgument, "guardrailIdentifier is required") } - if m.findGuardrailRecord(in.GuardrailIdentifier) == nil { + rec := m.findGuardrailRecord(in.GuardrailIdentifier) + if rec == nil { return nil, errors.Newf(errors.NotFound, "guardrail %q not found", in.GuardrailIdentifier) } + // A specific numbered version must exist; "" and "DRAFT" resolve to the + // working draft, which always exists. + if v := in.GuardrailVersion; v != "" && v != guardrailDraftVersion { + if _, ok := rec.versionSnapshot(v); !ok { + return nil, errors.Newf(errors.NotFound, "guardrail %q version %q not found", in.GuardrailIdentifier, v) + } + } + if in.Source != driver.GuardrailSourceInput && in.Source != driver.GuardrailSourceOutput { return nil, errors.Newf(errors.InvalidArgument, "invalid source %q: want INPUT or OUTPUT", in.Source) } diff --git a/providers/aws/bedrock/jobs.go b/providers/aws/bedrock/jobs.go index 7fec38b9..41cf5e12 100644 --- a/providers/aws/bedrock/jobs.go +++ b/providers/aws/bedrock/jobs.go @@ -184,9 +184,12 @@ func (m *Mock) CreateEvaluationJob(_ context.Context, cfg driver.EvaluationJobCo OutputDataS3URI: cfg.OutputDataS3URI, JobDescription: cfg.JobDescription, CustomerEncryptionKeyID: cfg.CustomerEncryptionKeyID, - Status: driver.JobCompleted, - CreationTime: now, - LastModifiedTime: now, + // Unlike import/copy jobs (which produce an artifact synchronously), + // evaluation is long-running, so it starts InProgress and stays there + // until StopEvaluationJob transitions it — making Stop a meaningful op. + Status: driver.JobInProgress, + CreationTime: now, + LastModifiedTime: now, } m.evalJobs.Set(cfg.JobName, job) m.setTags(jobARN, m.tagsFromMap(cfg.JobTags)) @@ -227,6 +230,13 @@ func (m *Mock) StopEvaluationJob(_ context.Context, jobIdentifier string) error return errors.Newf(errors.NotFound, "evaluation job %q not found", jobIdentifier) } + // Real AWS rejects stopping a job that is no longer in progress with a + // ConflictException. + if job.Status != driver.JobInProgress { + return errors.Newf(errors.FailedPrecondition, + "evaluation job %q is in terminal state %q and cannot be stopped", jobIdentifier, job.Status) + } + // Copy-on-write: never mutate the stored pointer in place, so concurrent // Get/List readers (which copy the stored value) can't race the write. updated := *job diff --git a/providers/aws/bedrock/marketplace_agreements.go b/providers/aws/bedrock/marketplace_agreements.go index 0e3299ad..db4b6623 100644 --- a/providers/aws/bedrock/marketplace_agreements.go +++ b/providers/aws/bedrock/marketplace_agreements.go @@ -136,14 +136,17 @@ func (m *Mock) RegisterMarketplaceModelEndpoint( return &result, nil } -// DeregisterMarketplaceModelEndpoint removes the Bedrock registration for an -// endpoint while leaving the endpoint record (and the underlying, unmodeled -// SageMaker endpoint) in place, so it can still be described or deleted. +// DeregisterMarketplaceModelEndpoint removes an endpoint's Bedrock registration. +// The endpoint is no longer tracked as a marketplace model endpoint, so a +// subsequent Get returns NotFound (the underlying, unmodeled SageMaker endpoint +// is unaffected), matching real AWS. func (m *Mock) DeregisterMarketplaceModelEndpoint(_ context.Context, endpointARN string) error { if !m.marketplaceEndpoints.Has(endpointARN) { return errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) } + m.marketplaceEndpoints.Delete(endpointARN) + return nil } diff --git a/providers/aws/bedrock/marketplace_agreements_test.go b/providers/aws/bedrock/marketplace_agreements_test.go index 0ee0d0a1..a1083afc 100644 --- a/providers/aws/bedrock/marketplace_agreements_test.go +++ b/providers/aws/bedrock/marketplace_agreements_test.go @@ -46,12 +46,23 @@ func TestMarketplaceEndpointLifecycle(t *testing.T) { assertEqual(t, bedrockdriver.MarketplaceEndpointStatusRegistered, reg.Status) assertEqual(t, "arn:model/source-2", reg.ModelSourceIdentifier) + // Deregister removes the Bedrock registration: a subsequent Get is NotFound. requireNoError(t, m.DeregisterMarketplaceModelEndpoint(ctx, endpoint.EndpointARN)) - requireNoError(t, m.DeleteMarketplaceModelEndpoint(ctx, endpoint.EndpointARN)) - _, err = m.GetMarketplaceModelEndpoint(ctx, endpoint.EndpointARN) assertError(t, err, true) + + // Delete also removes an endpoint (exercised on a fresh one). + ep2, err := m.CreateMarketplaceModelEndpoint(ctx, bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: "endpoint-2", + ModelSourceIdentifier: "arn:aws:sagemaker:us-east-1:aws:hub-content/model/2", + EndpointConfig: []byte(`{"sageMaker":{"instanceType":"ml.m5.large"}}`), + }) + requireNoError(t, err) + requireNoError(t, m.DeleteMarketplaceModelEndpoint(ctx, ep2.EndpointARN)) + + _, err = m.GetMarketplaceModelEndpoint(ctx, ep2.EndpointARN) + assertError(t, err, true) } func TestMarketplaceEndpointValidationAndErrors(t *testing.T) { diff --git a/providers/aws/bedrock/registries.go b/providers/aws/bedrock/registries.go index c546fdf0..88e5f643 100644 --- a/providers/aws/bedrock/registries.go +++ b/providers/aws/bedrock/registries.go @@ -48,14 +48,14 @@ func (m *Mock) CreateInferenceProfile(_ context.Context, cfg driver.InferencePro // GetInferenceProfile returns an inference profile by ID or ARN. func (m *Mock) GetInferenceProfile(_ context.Context, identifier string) (*driver.InferenceProfile, error) { if p, ok := m.inferenceProfiles.Get(identifier); ok { - result := *p + result := cloneInferenceProfile(p) return &result, nil } for _, p := range m.inferenceProfiles.All() { if p.ARN == identifier { - result := *p + result := cloneInferenceProfile(p) return &result, nil } @@ -70,7 +70,7 @@ func (m *Mock) ListInferenceProfiles(_ context.Context) ([]driver.InferenceProfi out := make([]driver.InferenceProfile, 0, len(all)) for _, p := range all { - out = append(out, *p) + out = append(out, cloneInferenceProfile(p)) } return out, nil @@ -150,7 +150,7 @@ func (m *Mock) GetPromptRouter(_ context.Context, promptRouterARN string) (*driv return nil, errors.Newf(errors.NotFound, "prompt router %q not found", promptRouterARN) } - result := *router + result := clonePromptRouter(router) return &result, nil } @@ -161,7 +161,7 @@ func (m *Mock) ListPromptRouters(_ context.Context) ([]driver.PromptRouter, erro out := make([]driver.PromptRouter, 0, len(all)) for _, router := range all { - out = append(out, *router) + out = append(out, clonePromptRouter(router)) } return out, nil @@ -225,7 +225,7 @@ func (m *Mock) GetAutomatedReasoningPolicy(_ context.Context, policyARN string) return nil, errors.Newf(errors.NotFound, "automated reasoning policy %q not found", policyARN) } - result := *policy + result := cloneARPolicy(policy) return &result, nil } @@ -236,7 +236,7 @@ func (m *Mock) ListAutomatedReasoningPolicies(_ context.Context) ([]driver.Autom out := make([]driver.AutomatedReasoningPolicy, 0, len(all)) for _, policy := range all { - out = append(out, *policy) + out = append(out, cloneARPolicy(policy)) } return out, nil @@ -295,3 +295,30 @@ func definitionHash(def []byte) string { return hex.EncodeToString(sum[:]) } + +// cloneInferenceProfile returns a value copy whose Models slice does not alias +// the stored profile, so callers can't mutate internal state via the result. +func cloneInferenceProfile(p *driver.InferenceProfile) driver.InferenceProfile { + out := *p + out.Models = append([]string(nil), p.Models...) + + return out +} + +// clonePromptRouter returns a value copy whose Models slice does not alias the +// stored router. +func clonePromptRouter(p *driver.PromptRouter) driver.PromptRouter { + out := *p + out.Models = append([]string(nil), p.Models...) + + return out +} + +// cloneARPolicy returns a value copy whose PolicyDefinition does not alias the +// stored policy. +func cloneARPolicy(p *driver.AutomatedReasoningPolicy) driver.AutomatedReasoningPolicy { + out := *p + out.PolicyDefinition = copyBytes(p.PolicyDefinition) + + return out +} diff --git a/server/aws/bedrock/errors.go b/server/aws/bedrock/errors.go index 87bc932f..5c57daab 100644 --- a/server/aws/bedrock/errors.go +++ b/server/aws/bedrock/errors.go @@ -35,7 +35,9 @@ func writeErr(w http.ResponseWriter, err error) { case cerrors.IsInvalidArgument(err): writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) case cerrors.IsFailedPrecondition(err): - writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + // A resource in a conflicting state (e.g. stopping a terminal job) maps + // to Bedrock's ConflictException, matching real AWS. + writeError(w, http.StatusConflict, "ConflictException", err.Error()) case cerrors.IsThrottled(err): writeError(w, http.StatusTooManyRequests, "ThrottlingException", err.Error()) default: diff --git a/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go b/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go index f78be1ab..c212d2ff 100644 --- a/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go +++ b/server/aws/bedrock/sdk_roundtrip_asyncjobs_test.go @@ -237,8 +237,8 @@ func TestSDKEvaluationJobLifecycle(t *testing.T) { t.Fatalf("GetEvaluationJob: %v", err) } - if got.Status != bedrocktypes.EvaluationJobStatusCompleted { - t.Fatalf("got status %q, want Completed", got.Status) + if got.Status != bedrocktypes.EvaluationJobStatusInProgress { + t.Fatalf("got status %q, want InProgress", got.Status) } if got.JobType != bedrocktypes.EvaluationJobTypeAutomated { diff --git a/server/aws/bedrock/sdk_roundtrip_marketplace_test.go b/server/aws/bedrock/sdk_roundtrip_marketplace_test.go index bfbcaf06..fc6e1dea 100644 --- a/server/aws/bedrock/sdk_roundtrip_marketplace_test.go +++ b/server/aws/bedrock/sdk_roundtrip_marketplace_test.go @@ -97,25 +97,41 @@ func TestSDKMarketplaceModelEndpointLifecycle(t *testing.T) { t.Fatalf("got status %q after register, want REGISTERED", reg.MarketplaceModelEndpoint.Status) } + // Deregister removes the Bedrock registration: a subsequent Get is NotFound. if _, err = client.DeregisterMarketplaceModelEndpoint(ctx, &awsbedrock.DeregisterMarketplaceModelEndpointInput{ EndpointArn: aws.String(arn), }); err != nil { t.Fatalf("DeregisterMarketplaceModelEndpoint: %v", err) } - if _, err = client.DeleteMarketplaceModelEndpoint(ctx, &awsbedrock.DeleteMarketplaceModelEndpointInput{ + _, err = client.GetMarketplaceModelEndpoint(ctx, &awsbedrock.GetMarketplaceModelEndpointInput{ EndpointArn: aws.String(arn), + }) + + var nfe *bedrocktypes.ResourceNotFoundException + if !errors.As(err, &nfe) { + t.Fatalf("expected ResourceNotFoundException after deregister, got %T: %v", err, err) + } + + // Delete also removes an endpoint (exercised over the SDK on a fresh one). + create2, err := client.CreateMarketplaceModelEndpoint(ctx, &awsbedrock.CreateMarketplaceModelEndpointInput{ + EndpointConfig: cfg, EndpointName: aws.String("endpoint-2"), ModelSourceIdentifier: aws.String(source), AcceptEula: true, + }) + if err != nil { + t.Fatalf("CreateMarketplaceModelEndpoint(2): %v", err) + } + + if _, err = client.DeleteMarketplaceModelEndpoint(ctx, &awsbedrock.DeleteMarketplaceModelEndpointInput{ + EndpointArn: create2.MarketplaceModelEndpoint.EndpointArn, }); err != nil { t.Fatalf("DeleteMarketplaceModelEndpoint: %v", err) } _, err = client.GetMarketplaceModelEndpoint(ctx, &awsbedrock.GetMarketplaceModelEndpointInput{ - EndpointArn: aws.String(arn), + EndpointArn: create2.MarketplaceModelEndpoint.EndpointArn, }) - - var nfe *bedrocktypes.ResourceNotFoundException if !errors.As(err, &nfe) { - t.Fatalf("expected ResourceNotFoundException, got %T: %v", err, err) + t.Fatalf("expected ResourceNotFoundException after delete, got %T: %v", err, err) } } diff --git a/services/bedrockagent/bedrockagent_test.go b/services/bedrockagent/bedrockagent_test.go new file mode 100644 index 00000000..1608c11c --- /dev/null +++ b/services/bedrockagent/bedrockagent_test.go @@ -0,0 +1,99 @@ +package bedrockagent + +import ( + "context" + "testing" + "time" + + "github.com/stackshy/cloudemu/v2/config" + provider "github.com/stackshy/cloudemu/v2/providers/aws/bedrockagent" + "github.com/stackshy/cloudemu/v2/services/bedrockagent/driver" +) + +func newService() *BedrockAgent { + fc := config.NewFakeClock(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)) + opts := config.NewOptions(config.WithClock(fc), config.WithRegion("us-east-1")) + + return NewBedrockAgent(provider.New(opts)) +} + +func TestServiceAgentLifecycle(t *testing.T) { + svc := newService() + ctx := context.Background() + + agent, err := svc.CreateAgent(ctx, driver.AgentConfig{Name: "svc-agent", FoundationModel: "anthropic.claude-3-sonnet-20240229-v1:0"}) + if err != nil { + t.Fatalf("CreateAgent: %v", err) + } + + if agent.ID == "" { + t.Fatal("expected an agent id") + } + + got, err := svc.GetAgent(ctx, agent.ID) + if err != nil { + t.Fatalf("GetAgent: %v", err) + } + + if got.Name != "svc-agent" { + t.Fatalf("got name %q, want svc-agent", got.Name) + } + + prepared, err := svc.PrepareAgent(ctx, agent.ID) + if err != nil { + t.Fatalf("PrepareAgent: %v", err) + } + + if prepared.Status != driver.AgentPrepared { + t.Fatalf("got status %q, want %q", prepared.Status, driver.AgentPrepared) + } + + agents, err := svc.ListAgents(ctx) + if err != nil { + t.Fatalf("ListAgents: %v", err) + } + + if len(agents) != 1 { + t.Fatalf("got %d agents, want 1", len(agents)) + } + + if _, err := svc.DeleteAgent(ctx, agent.ID); err != nil { + t.Fatalf("DeleteAgent: %v", err) + } + + if _, err := svc.GetAgent(ctx, agent.ID); err == nil { + t.Fatal("expected GetAgent to fail after delete") + } +} + +func TestServiceKnowledgeBaseLifecycle(t *testing.T) { + svc := newService() + ctx := context.Background() + + kb, err := svc.CreateKnowledgeBase(ctx, driver.KnowledgeBaseConfig{ + Name: "svc-kb", + RoleArn: "arn:aws:iam::123456789012:role/r", + KnowledgeBaseConfiguration: []byte(`{"type":"VECTOR"}`), + }) + if err != nil { + t.Fatalf("CreateKnowledgeBase: %v", err) + } + + got, err := svc.GetKnowledgeBase(ctx, kb.ID) + if err != nil { + t.Fatalf("GetKnowledgeBase: %v", err) + } + + if got.Name != "svc-kb" { + t.Fatalf("got name %q, want svc-kb", got.Name) + } + + list, err := svc.ListKnowledgeBases(ctx) + if err != nil { + t.Fatalf("ListKnowledgeBases: %v", err) + } + + if len(list) != 1 { + t.Fatalf("got %d knowledge bases, want 1", len(list)) + } +} From a84241182167cab084780f26e502c4e0161e3074 Mon Sep 17 00:00:00 2001 From: Satyam Trivedi Date: Wed, 29 Jul 2026 12:36:50 +0530 Subject: [PATCH 4/5] =?UTF-8?q?fix(bedrock):=20address=20in-depth=20review?= =?UTF-8?q?=20=E2=80=94=20cascade=20delete,=20copy-out,=20error=20taxonomy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the fourth PR review on #214. - bedrock-agent cascade delete: DeleteKnowledgeBase now also removes its data sources and ingestion jobs; DeleteDataSource removes its ingestion jobs; DeleteAgent removes its aliases. No more orphaned children in the store. - Copy-out immutability on read: Get/List/Update/Create-return paths now clone json.RawMessage / []byte fields out (knowledge-base + data-source + flow + prompt configs; marketplace endpointConfig; evaluation-job evaluation/inference configs), so a caller mutating a returned value can't corrupt stored state — consistent with the guardrail/AR-policy paths. - Error taxonomy: writeErr in bedrock, bedrock-agent, and bedrock-agent-runtime now maps PermissionDenied -> 403 AccessDeniedException and ResourceExhausted -> 400 ServiceQuotaExceededException instead of falling through to 500. - CI: added a scoped `go test -race` step for the Bedrock packages so their concurrency tests actually exercise the race detector in CI. - Docs: recorded that bedrock-agent versioning/aliases-beyond-create, action groups, and collaborators are out of scope for this iteration, and that jobs complete synchronously / inference responses are deterministic simulations. Verified: go build/vet/gofmt/golangci-lint clean, go test ./... (204 pkgs) passing, and go test -race on the bedrock packages clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/ci.yml | 4 + docs/sdk-server.md | 6 +- .../aws/bedrock/asyncinvoke_jobs_test.go | 28 +++++ providers/aws/bedrock/jobs.go | 15 ++- .../aws/bedrock/marketplace_agreements.go | 20 +++- .../bedrock/marketplace_agreements_test.go | 26 ++++ providers/aws/bedrockagent/agents.go | 14 ++- .../aws/bedrockagent/bedrockagent_test.go | 112 ++++++++++++++++++ providers/aws/bedrockagent/datasources.go | 49 +++++++- providers/aws/bedrockagent/flows.go | 19 ++- providers/aws/bedrockagent/knowledgebases.go | 23 +++- providers/aws/bedrockagent/prompts.go | 17 ++- server/aws/bedrock/errors.go | 5 + server/aws/bedrockagent/errors.go | 4 + server/aws/bedrockagentruntime/errors.go | 4 + 15 files changed, 317 insertions(+), 29 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3b02a5d7..41083f36 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -88,6 +88,10 @@ jobs: restore-keys: | go-${{ runner.os }}- - run: go test ./... + # Run the Bedrock packages under the race detector so their concurrency + # tests (guardrail versions, evaluation-job stop, etc.) actually exercise + # -race in CI rather than being inert. + - run: go test -race ./providers/aws/bedrock/... ./server/aws/bedrock/... ./providers/aws/bedrockagent/... ./server/aws/bedrockagent/... ./server/aws/bedrockagentruntime/... tidy: name: Tidy diff --git a/docs/sdk-server.md b/docs/sdk-server.md index c362899a..356c3d40 100644 --- a/docs/sdk-server.md +++ b/docs/sdk-server.md @@ -176,7 +176,7 @@ Region, credentials, and tokens can be any dummy values — the server doesn't v | **Resource Explorer 2** *(JSON)* | Search — free-text plus filter expression over the cross-service inventory; results include ARN, resource type, region, owning account, and tags | | **Resource Groups Tagging API** *(JSON-RPC)* | GetResources (filter by `ResourceTypeFilters` + `TagFilters`, paginated), TagResources, UntagResources, GetTagKeys, GetTagValues | | **Bedrock** *(REST + JSON, `bedrock` + `bedrock-runtime`)* | Control plane: foundation models (List/Get), model-customization jobs (Create/Get/List), custom models (List/Get/Delete), Guardrails (Create/Get/List/Update/Delete + CreateGuardrailVersion, with topic/content/word/sensitive-info/contextual-grounding policy configs and version snapshots), Provisioned Throughput (Create/Get/List/Delete), invocation-logging config (Put/Get/Delete), resource tagging (Tag/Untag/ListTagsForResource), model import jobs, model copy jobs, evaluation jobs (Create/Get/List/Stop), inference profiles (Create/Get/List/Delete), prompt routers (Create/Get/List/Delete), marketplace model endpoints (Create/Get/List/Update/Delete/Register/Deregister), foundation-model agreements (Create/Delete/ListOffers/GetAvailability), automated-reasoning policies (Create/Get/List/Update/Delete). Runtime: InvokeModel (family-aware response envelopes), Converse, ConverseStream + InvokeModelWithResponseStream (eventstream), CountTokens, ApplyGuardrail, and async invoke (Start/Get/List). | -| **Bedrock Agent** *(REST + JSON, `bedrock-agent` + `bedrock-agent-runtime`)* | Control plane: agents (Create/Get/List/Update/Delete/Prepare + alias), knowledge bases (CRUD), data sources (CRUD + StartIngestionJob), flows (CRUD + Prepare), prompts (CRUD). Runtime: InvokeAgent (eventstream), Retrieve, RetrieveAndGenerate. | +| **Bedrock Agent** *(REST + JSON, `bedrock-agent` + `bedrock-agent-runtime`)* | Control plane: agents (Create/Get/List/Update/Delete/Prepare + alias), knowledge bases (CRUD), data sources (CRUD + StartIngestionJob), flows (CRUD + Prepare), prompts (CRUD). Runtime: InvokeAgent (eventstream), Retrieve, RetrieveAndGenerate. Scope: the core resource lifecycle and runtime data plane above; **agent versioning/aliases beyond basic create, action groups, and agent collaborators are out of scope** for this iteration. | ### Azure (`server/azure/`) @@ -323,7 +323,9 @@ Kubernetes ships as **two cooperating handlers**: per-provider control planes (E The data plane intentionally has no controllers — Deployments don't spawn ReplicaSets, Pods stay Pending, Endpoints are empty stubs. RBAC, subresources, PV/PVC, StatefulSet/DaemonSet/Job/CronJob, and Ingress are out of scope. -Two provider-specific services also ship as full SDK-compat handlers. **AWS Bedrock** covers the `bedrock` control plane (foundation models, customization jobs, custom models, guardrails with policy configs + versions, provisioned throughput, invocation logging, resource tagging, model import/copy/evaluation jobs, inference profiles, prompt routers, marketplace model endpoints, foundation-model agreements, and automated-reasoning policies) and the `bedrock-runtime` data plane (InvokeModel with family-aware response envelopes, Converse, streaming ConverseStream / InvokeModelWithResponseStream over `vnd.amazon.eventstream`, CountTokens, ApplyGuardrail, and async invoke). A companion **AWS Bedrock Agent** handler covers the `bedrock-agent` control plane (agents, knowledge bases, data sources, flows, prompts) and the `bedrock-agent-runtime` data plane (InvokeAgent streaming, Retrieve, RetrieveAndGenerate); its runtime handler registers before the control plane and matches only POST so the two never collide on the shared `/agents` and `/knowledgebases` roots. **Azure Databricks** covers the `armdatabricks` ARM workspace resource plus the `databricks-sdk-go` workspace data plane — clusters, instance pools, jobs and runs, cluster policies, libraries, permissions, secrets, tokens, git credentials, repos, DBFS, workspace notebooks/directories, SQL warehouses, pipelines, serving endpoints, SCIM identity, and Unity Catalog. +Two provider-specific services also ship as full SDK-compat handlers. **AWS Bedrock** covers the `bedrock` control plane (foundation models, customization jobs, custom models, guardrails with policy configs + versions, provisioned throughput, invocation logging, resource tagging, model import/copy/evaluation jobs, inference profiles, prompt routers, marketplace model endpoints, foundation-model agreements, and automated-reasoning policies) and the `bedrock-runtime` data plane (InvokeModel with family-aware response envelopes, Converse, streaming ConverseStream / InvokeModelWithResponseStream over `vnd.amazon.eventstream`, CountTokens, ApplyGuardrail, and async invoke). A companion **AWS Bedrock Agent** handler covers the `bedrock-agent` control plane (agents, knowledge bases, data sources, flows, prompts) and the `bedrock-agent-runtime` data plane (InvokeAgent streaming, Retrieve, RetrieveAndGenerate); its runtime handler registers before the control plane and matches only POST so the two never collide on the shared `/agents` and `/knowledgebases` roots. `bedrock-agent` coverage is intentionally scoped to this core resource lifecycle and runtime data plane — agent versioning/aliases beyond basic create, action groups, and agent collaborators are out of scope for this iteration. **Azure Databricks** covers the `armdatabricks` ARM workspace resource plus the `databricks-sdk-go` workspace data plane — clusters, instance pools, jobs and runs, cluster policies, libraries, permissions, secrets, tokens, git credentials, repos, DBFS, workspace notebooks/directories, SQL warehouses, pipelines, serving endpoints, SCIM identity, and Unity Catalog. + +**Emulation caveats (Bedrock).** Long-running jobs — model customization, import, and copy jobs (evaluation jobs start `InProgress`) — complete synchronously in the emulator, so Get/List observe a terminal state immediately rather than polling through intermediate progress. Inference and agent responses (InvokeModel, Converse, InvokeAgent, RetrieveAndGenerate) are deterministic simulations, not real model output. The remaining service domains (DNS, Load Balancer, Cache, Secrets, Logging, Notifications, Container Registry, Event Bus) have full driver implementations in `providers/{aws,azure,gcp}/`; SDK-compat handlers are added in lockstep across all 3 providers as each domain ships. diff --git a/providers/aws/bedrock/asyncinvoke_jobs_test.go b/providers/aws/bedrock/asyncinvoke_jobs_test.go index 46af50b4..52bf7a30 100644 --- a/providers/aws/bedrock/asyncinvoke_jobs_test.go +++ b/providers/aws/bedrock/asyncinvoke_jobs_test.go @@ -166,6 +166,34 @@ func TestEvaluationJobLifecycle(t *testing.T) { assertError(t, m.StopEvaluationJob(ctx, "missing"), true) } +// TestEvaluationJobCopyOut verifies that mutating the EvaluationConfig bytes of +// a returned job does not affect the stored value. +func TestEvaluationJobCopyOut(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + _, err := m.CreateEvaluationJob(ctx, bedrockdriver.EvaluationJobConfig{ + JobName: "eval-copyout", + RoleARN: "arn:aws:iam::123456789012:role/bedrock", + EvaluationConfig: []byte(`{"automated":{}}`), + InferenceConfig: []byte(`{"models":[]}`), + OutputDataS3URI: "s3://bucket/eval/", + }) + requireNoError(t, err) + + got, err := m.GetEvaluationJob(ctx, "eval-copyout") + requireNoError(t, err) + + // Mutate the returned bytes; the store must be unaffected. + for i := range got.EvaluationConfig { + got.EvaluationConfig[i] = 'X' + } + + again, err := m.GetEvaluationJob(ctx, "eval-copyout") + requireNoError(t, err) + assertEqual(t, `{"automated":{}}`, string(again.EvaluationConfig)) +} + func TestEvaluationJobTypeHuman(t *testing.T) { m := newTestMock() diff --git a/providers/aws/bedrock/jobs.go b/providers/aws/bedrock/jobs.go index 41cf5e12..dd940ed6 100644 --- a/providers/aws/bedrock/jobs.go +++ b/providers/aws/bedrock/jobs.go @@ -206,7 +206,7 @@ func (m *Mock) GetEvaluationJob(_ context.Context, jobIdentifier string) (*drive return nil, errors.Newf(errors.NotFound, "evaluation job %q not found", jobIdentifier) } - result := *job + result := cloneEvaluationJob(job) return &result, nil } @@ -217,7 +217,7 @@ func (m *Mock) ListEvaluationJobs(_ context.Context) ([]driver.EvaluationJob, er out := make([]driver.EvaluationJob, 0, len(all)) for _, job := range all { - out = append(out, *job) + out = append(out, cloneEvaluationJob(job)) } return out, nil @@ -261,6 +261,17 @@ func (m *Mock) findEvalJob(id string) *driver.EvaluationJob { return nil } +// cloneEvaluationJob returns a value copy whose EvaluationConfig and +// InferenceConfig do not alias the stored job, so callers can't mutate internal +// state via the result. +func cloneEvaluationJob(j *driver.EvaluationJob) driver.EvaluationJob { + out := *j + out.EvaluationConfig = copyBytes(j.EvaluationConfig) + out.InferenceConfig = copyBytes(j.InferenceConfig) + + return out +} + // evaluationJobType derives the job type from the evaluationConfig document: a // "human" member yields Human, otherwise Automated. func evaluationJobType(cfg []byte) string { diff --git a/providers/aws/bedrock/marketplace_agreements.go b/providers/aws/bedrock/marketplace_agreements.go index db4b6623..93b4c0b9 100644 --- a/providers/aws/bedrock/marketplace_agreements.go +++ b/providers/aws/bedrock/marketplace_agreements.go @@ -46,7 +46,7 @@ func (m *Mock) CreateMarketplaceModelEndpoint( m.marketplaceEndpoints.Set(arn, endpoint) m.setTags(arn, m.tagsFromMap(cfg.Tags)) - result := *endpoint + result := cloneMarketplaceEndpoint(endpoint) return &result, nil } @@ -58,7 +58,7 @@ func (m *Mock) GetMarketplaceModelEndpoint(_ context.Context, endpointARN string return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointARN) } - result := *endpoint + result := cloneMarketplaceEndpoint(endpoint) return &result, nil } @@ -69,7 +69,7 @@ func (m *Mock) ListMarketplaceModelEndpoints(_ context.Context) ([]driver.Market out := make([]driver.MarketplaceEndpoint, 0, len(all)) for _, endpoint := range all { - out = append(out, *endpoint) + out = append(out, cloneMarketplaceEndpoint(endpoint)) } return out, nil @@ -94,7 +94,7 @@ func (m *Mock) UpdateMarketplaceModelEndpoint( updated.UpdatedAt = m.now() m.marketplaceEndpoints.Set(endpointARN, &updated) - result := updated + result := cloneMarketplaceEndpoint(&updated) return &result, nil } @@ -131,7 +131,7 @@ func (m *Mock) RegisterMarketplaceModelEndpoint( updated.UpdatedAt = m.now() m.marketplaceEndpoints.Set(endpointIdentifier, &updated) - result := updated + result := cloneMarketplaceEndpoint(&updated) return &result, nil } @@ -213,6 +213,16 @@ func (m *Mock) GetFoundationModelAvailability( return out, nil } +// cloneMarketplaceEndpoint returns a value copy whose EndpointConfig does not +// alias the stored endpoint, so callers can't mutate internal state via the +// result. +func cloneMarketplaceEndpoint(e *driver.MarketplaceEndpoint) driver.MarketplaceEndpoint { + out := *e + out.EndpointConfig = copyBytes(e.EndpointConfig) + + return out +} + // copyBytes returns a copy of b so stored payloads never alias caller memory. func copyBytes(b []byte) []byte { if len(b) == 0 { diff --git a/providers/aws/bedrock/marketplace_agreements_test.go b/providers/aws/bedrock/marketplace_agreements_test.go index a1083afc..b84cab2f 100644 --- a/providers/aws/bedrock/marketplace_agreements_test.go +++ b/providers/aws/bedrock/marketplace_agreements_test.go @@ -97,6 +97,32 @@ func TestMarketplaceEndpointValidationAndErrors(t *testing.T) { assertError(t, m.DeleteMarketplaceModelEndpoint(ctx, "arn:missing"), true) } +// TestMarketplaceEndpointCopyOut verifies that mutating the EndpointConfig +// bytes of a returned endpoint does not affect the stored value. +func TestMarketplaceEndpointCopyOut(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + created, err := m.CreateMarketplaceModelEndpoint(ctx, bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: "endpoint-copyout", + ModelSourceIdentifier: "arn:aws:sagemaker:us-east-1:aws:hub-content/model/1", + EndpointConfig: []byte(`{"sageMaker":{"instanceType":"ml.m5.large"}}`), + }) + requireNoError(t, err) + + got, err := m.GetMarketplaceModelEndpoint(ctx, created.EndpointARN) + requireNoError(t, err) + + // Mutate the returned bytes; the store must be unaffected. + for i := range got.EndpointConfig { + got.EndpointConfig[i] = 'X' + } + + again, err := m.GetMarketplaceModelEndpoint(ctx, created.EndpointARN) + requireNoError(t, err) + assertEqual(t, `{"sageMaker":{"instanceType":"ml.m5.large"}}`, string(again.EndpointConfig)) +} + func TestFoundationModelAgreementLifecycle(t *testing.T) { m := newTestMock() ctx := context.Background() diff --git a/providers/aws/bedrockagent/agents.go b/providers/aws/bedrockagent/agents.go index 153f580f..404dfa7b 100644 --- a/providers/aws/bedrockagent/agents.go +++ b/providers/aws/bedrockagent/agents.go @@ -98,17 +98,29 @@ func (m *Mock) UpdateAgent(_ context.Context, agentID string, cfg driver.AgentCo return &result, nil } -// DeleteAgent deletes an agent and returns its terminal status. +// DeleteAgent deletes an agent and, cascading like real AWS, every alias that +// belongs to it. func (m *Mock) DeleteAgent(_ context.Context, agentID string) (string, error) { if !m.agents.Has(agentID) { return "", errors.Newf(errors.NotFound, "agent %q not found", agentID) } m.agents.Delete(agentID) + m.deleteAliasesForAgent(agentID) return statusDeleting, nil } +// deleteAliasesForAgent removes every alias belonging to agentID. All() returns +// a snapshot, so deleting while ranging is safe. +func (m *Mock) deleteAliasesForAgent(agentID string) { + for id, alias := range m.aliases.All() { + if alias.AgentID == agentID { + m.aliases.Delete(id) + } + } +} + // PrepareAgent prepares an agent, transitioning it to PREPARED. func (m *Mock) PrepareAgent(_ context.Context, agentID string) (*driver.Agent, error) { agent, ok := m.agents.Get(agentID) diff --git a/providers/aws/bedrockagent/bedrockagent_test.go b/providers/aws/bedrockagent/bedrockagent_test.go index 3607ca02..dc3ae957 100644 --- a/providers/aws/bedrockagent/bedrockagent_test.go +++ b/providers/aws/bedrockagent/bedrockagent_test.go @@ -100,6 +100,118 @@ func TestKnowledgeBaseAndDataSource(t *testing.T) { assert.True(t, cerrors.IsNotFound(err)) } +func TestDeleteKnowledgeBaseCascade(t *testing.T) { + m := newMock() + ctx := context.Background() + + kb, err := m.CreateKnowledgeBase(ctx, driver.KnowledgeBaseConfig{ + Name: "kb1", + RoleArn: "role", + KnowledgeBaseConfiguration: json.RawMessage(`{"type":"VECTOR"}`), + }) + require.NoError(t, err) + + ds, err := m.CreateDataSource(ctx, driver.DataSourceConfig{ + KnowledgeBaseID: kb.ID, + Name: "ds1", + DataSourceConfiguration: json.RawMessage(`{"type":"S3"}`), + }) + require.NoError(t, err) + + job, err := m.StartIngestionJob(ctx, kb.ID, ds.ID, "reindex") + require.NoError(t, err) + + _, err = m.DeleteKnowledgeBase(ctx, kb.ID) + require.NoError(t, err) + + _, err = m.GetDataSource(ctx, kb.ID, ds.ID) + assert.True(t, cerrors.IsNotFound(err), "data source should be cascade-deleted") + + sources, err := m.ListDataSources(ctx, kb.ID) + require.NoError(t, err) + assert.Empty(t, sources) + + assert.False(t, m.jobs.Has(job.ID), "ingestion job should be cascade-deleted") +} + +func TestDeleteDataSourceCascadesJobs(t *testing.T) { + m := newMock() + ctx := context.Background() + + kb, err := m.CreateKnowledgeBase(ctx, driver.KnowledgeBaseConfig{ + Name: "kb1", + RoleArn: "role", + KnowledgeBaseConfiguration: json.RawMessage(`{"type":"VECTOR"}`), + }) + require.NoError(t, err) + + ds, err := m.CreateDataSource(ctx, driver.DataSourceConfig{ + KnowledgeBaseID: kb.ID, + Name: "ds1", + DataSourceConfiguration: json.RawMessage(`{"type":"S3"}`), + }) + require.NoError(t, err) + + job, err := m.StartIngestionJob(ctx, kb.ID, ds.ID, "reindex") + require.NoError(t, err) + + _, err = m.DeleteDataSource(ctx, kb.ID, ds.ID) + require.NoError(t, err) + + assert.False(t, m.jobs.Has(job.ID), "ingestion job should be cascade-deleted") + // The knowledge base itself must survive its data source's deletion. + assert.True(t, m.knowledge.Has(kb.ID)) +} + +func TestDeleteAgentCascadesAliases(t *testing.T) { + m := newMock() + ctx := context.Background() + + agent, err := m.CreateAgent(ctx, driver.AgentConfig{Name: "a1", FoundationModel: "fm"}) + require.NoError(t, err) + + alias, err := m.CreateAgentAlias(ctx, driver.AgentAliasConfig{AgentID: agent.ID, Name: "prod"}) + require.NoError(t, err) + require.True(t, m.aliases.Has(alias.ID)) + + _, err = m.DeleteAgent(ctx, agent.ID) + require.NoError(t, err) + + assert.False(t, m.aliases.Has(alias.ID), "alias should be cascade-deleted") +} + +func TestDataSourceCopyOutImmutable(t *testing.T) { + m := newMock() + ctx := context.Background() + + kb, err := m.CreateKnowledgeBase(ctx, driver.KnowledgeBaseConfig{ + Name: "kb1", + RoleArn: "role", + KnowledgeBaseConfiguration: json.RawMessage(`{"type":"VECTOR"}`), + }) + require.NoError(t, err) + + original := `{"type":"S3"}` + ds, err := m.CreateDataSource(ctx, driver.DataSourceConfig{ + KnowledgeBaseID: kb.ID, + Name: "ds1", + DataSourceConfiguration: json.RawMessage(original), + }) + require.NoError(t, err) + + got, err := m.GetDataSource(ctx, kb.ID, ds.ID) + require.NoError(t, err) + + // Mutate the returned config bytes; stored state must not change. + for i := range got.DataSourceConfiguration { + got.DataSourceConfiguration[i] = 'X' + } + + again, err := m.GetDataSource(ctx, kb.ID, ds.ID) + require.NoError(t, err) + assert.JSONEq(t, original, string(again.DataSourceConfiguration)) +} + func TestDataSourceRequiresKnowledgeBase(t *testing.T) { m := newMock() diff --git a/providers/aws/bedrockagent/datasources.go b/providers/aws/bedrockagent/datasources.go index 9faa05a0..aed3ccec 100644 --- a/providers/aws/bedrockagent/datasources.go +++ b/providers/aws/bedrockagent/datasources.go @@ -39,7 +39,7 @@ func (m *Mock) CreateDataSource(_ context.Context, cfg driver.DataSourceConfig) } m.dataSource.Set(id, ds) - result := *ds + result := cloneDataSource(ds) return &result, nil } @@ -51,7 +51,7 @@ func (m *Mock) GetDataSource(_ context.Context, kbID, dsID string) (*driver.Data return nil, errors.Newf(errors.NotFound, "data source %q not found", dsID) } - result := *ds + result := cloneDataSource(ds) return &result, nil } @@ -63,7 +63,7 @@ func (m *Mock) ListDataSources(_ context.Context, kbID string) ([]driver.DataSou for _, ds := range all { if ds.KnowledgeBaseID == kbID { - out = append(out, *ds) + out = append(out, cloneDataSource(ds)) } } @@ -91,18 +91,20 @@ func (m *Mock) UpdateDataSource(_ context.Context, cfg driver.DataSourceConfig, m.dataSource.Set(dsID, &updated) - result := updated + result := cloneDataSource(&updated) return &result, nil } -// DeleteDataSource deletes a data source and returns its terminal status. +// DeleteDataSource deletes a data source and, cascading like real AWS, every +// ingestion job that belongs to it. func (m *Mock) DeleteDataSource(_ context.Context, kbID, dsID string) (string, error) { if m.findDataSource(kbID, dsID) == nil { return "", errors.Newf(errors.NotFound, "data source %q not found", dsID) } m.dataSource.Delete(dsID) + m.deleteJobsForDataSource(dsID) return statusDeleting, nil } @@ -140,3 +142,40 @@ func (m *Mock) findDataSource(kbID, dsID string) *driver.DataSource { return ds } + +// deleteDataSourcesForKnowledgeBase removes every data source belonging to kbID. +// All() returns a snapshot, so deleting while ranging is safe. +func (m *Mock) deleteDataSourcesForKnowledgeBase(kbID string) { + for id, ds := range m.dataSource.All() { + if ds.KnowledgeBaseID == kbID { + m.dataSource.Delete(id) + } + } +} + +// deleteJobsForKnowledgeBase removes every ingestion job belonging to kbID. +func (m *Mock) deleteJobsForKnowledgeBase(kbID string) { + for id, job := range m.jobs.All() { + if job.KnowledgeBaseID == kbID { + m.jobs.Delete(id) + } + } +} + +// deleteJobsForDataSource removes every ingestion job belonging to dsID. +func (m *Mock) deleteJobsForDataSource(dsID string) { + for id, job := range m.jobs.All() { + if job.DataSourceID == dsID { + m.jobs.Delete(id) + } + } +} + +// cloneDataSource returns a value copy whose RawMessage config field does not +// alias the stored data source, so callers can't mutate internal state. +func cloneDataSource(ds *driver.DataSource) driver.DataSource { + out := *ds + out.DataSourceConfiguration = copyRaw(ds.DataSourceConfiguration) + + return out +} diff --git a/providers/aws/bedrockagent/flows.go b/providers/aws/bedrockagent/flows.go index 170a94b9..8a797d04 100644 --- a/providers/aws/bedrockagent/flows.go +++ b/providers/aws/bedrockagent/flows.go @@ -36,7 +36,7 @@ func (m *Mock) CreateFlow(_ context.Context, cfg driver.FlowConfig) (*driver.Flo } m.flows.Set(id, flow) - result := *flow + result := cloneFlow(flow) return &result, nil } @@ -48,7 +48,7 @@ func (m *Mock) GetFlow(_ context.Context, id string) (*driver.Flow, error) { return nil, errors.Newf(errors.NotFound, "flow %q not found", id) } - result := *flow + result := cloneFlow(flow) return &result, nil } @@ -59,7 +59,7 @@ func (m *Mock) ListFlows(_ context.Context) ([]driver.Flow, error) { out := make([]driver.Flow, 0, len(all)) for _, f := range all { - out = append(out, *f) + out = append(out, cloneFlow(f)) } return out, nil @@ -87,7 +87,7 @@ func (m *Mock) UpdateFlow(_ context.Context, id string, cfg driver.FlowConfig) ( m.flows.Set(id, &updated) - result := updated + result := cloneFlow(&updated) return &result, nil } @@ -115,7 +115,16 @@ func (m *Mock) PrepareFlow(_ context.Context, id string) (*driver.Flow, error) { updated.UpdatedAt = m.now() m.flows.Set(id, &updated) - result := updated + result := cloneFlow(&updated) return &result, nil } + +// cloneFlow returns a value copy whose Definition does not alias the stored +// flow, so callers can't mutate internal state via the result. +func cloneFlow(f *driver.Flow) driver.Flow { + out := *f + out.Definition = copyRaw(f.Definition) + + return out +} diff --git a/providers/aws/bedrockagent/knowledgebases.go b/providers/aws/bedrockagent/knowledgebases.go index 9b68d59f..64b1644d 100644 --- a/providers/aws/bedrockagent/knowledgebases.go +++ b/providers/aws/bedrockagent/knowledgebases.go @@ -37,7 +37,7 @@ func (m *Mock) CreateKnowledgeBase(_ context.Context, cfg driver.KnowledgeBaseCo } m.knowledge.Set(id, kb) - result := *kb + result := cloneKnowledgeBase(kb) return &result, nil } @@ -49,7 +49,7 @@ func (m *Mock) GetKnowledgeBase(_ context.Context, id string) (*driver.Knowledge return nil, errors.Newf(errors.NotFound, "knowledge base %q not found", id) } - result := *kb + result := cloneKnowledgeBase(kb) return &result, nil } @@ -60,7 +60,7 @@ func (m *Mock) ListKnowledgeBases(_ context.Context) ([]driver.KnowledgeBase, er out := make([]driver.KnowledgeBase, 0, len(all)) for _, kb := range all { - out = append(out, *kb) + out = append(out, cloneKnowledgeBase(kb)) } return out, nil @@ -91,18 +91,31 @@ func (m *Mock) UpdateKnowledgeBase(_ context.Context, id string, cfg driver.Know m.knowledge.Set(id, &updated) - result := updated + result := cloneKnowledgeBase(&updated) return &result, nil } -// DeleteKnowledgeBase deletes a knowledge base and returns its terminal status. +// DeleteKnowledgeBase deletes a knowledge base and, cascading like real AWS, +// every data source and ingestion job that belongs to it. func (m *Mock) DeleteKnowledgeBase(_ context.Context, id string) (string, error) { if !m.knowledge.Has(id) { return "", errors.Newf(errors.NotFound, "knowledge base %q not found", id) } m.knowledge.Delete(id) + m.deleteDataSourcesForKnowledgeBase(id) + m.deleteJobsForKnowledgeBase(id) return statusDeleting, nil } + +// cloneKnowledgeBase returns a value copy whose RawMessage config fields do not +// alias the stored knowledge base, so callers can't mutate internal state. +func cloneKnowledgeBase(kb *driver.KnowledgeBase) driver.KnowledgeBase { + out := *kb + out.KnowledgeBaseConfiguration = copyRaw(kb.KnowledgeBaseConfiguration) + out.StorageConfiguration = copyRaw(kb.StorageConfiguration) + + return out +} diff --git a/providers/aws/bedrockagent/prompts.go b/providers/aws/bedrockagent/prompts.go index 4f27e0d0..c3b83765 100644 --- a/providers/aws/bedrockagent/prompts.go +++ b/providers/aws/bedrockagent/prompts.go @@ -32,7 +32,7 @@ func (m *Mock) CreatePrompt(_ context.Context, cfg driver.PromptConfig) (*driver } m.prompts.Set(id, prompt) - result := *prompt + result := clonePrompt(prompt) return &result, nil } @@ -44,7 +44,7 @@ func (m *Mock) GetPrompt(_ context.Context, id string) (*driver.Prompt, error) { return nil, errors.Newf(errors.NotFound, "prompt %q not found", id) } - result := *prompt + result := clonePrompt(prompt) return &result, nil } @@ -55,7 +55,7 @@ func (m *Mock) ListPrompts(_ context.Context) ([]driver.Prompt, error) { out := make([]driver.Prompt, 0, len(all)) for _, p := range all { - out = append(out, *p) + out = append(out, clonePrompt(p)) } return out, nil @@ -82,7 +82,7 @@ func (m *Mock) UpdatePrompt(_ context.Context, id string, cfg driver.PromptConfi m.prompts.Set(id, &updated) - result := updated + result := clonePrompt(&updated) return &result, nil } @@ -97,3 +97,12 @@ func (m *Mock) DeletePrompt(_ context.Context, id string) (string, error) { return id, nil } + +// clonePrompt returns a value copy whose Variants do not alias the stored +// prompt, so callers can't mutate internal state via the result. +func clonePrompt(p *driver.Prompt) driver.Prompt { + out := *p + out.Variants = copyRaw(p.Variants) + + return out +} diff --git a/server/aws/bedrock/errors.go b/server/aws/bedrock/errors.go index 5c57daab..6b5d58fe 100644 --- a/server/aws/bedrock/errors.go +++ b/server/aws/bedrock/errors.go @@ -34,6 +34,11 @@ func writeErr(w http.ResponseWriter, err error) { writeError(w, http.StatusConflict, "ConflictException", err.Error()) case cerrors.IsInvalidArgument(err): writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) + case cerrors.IsPermissionDenied(err): + writeError(w, http.StatusForbidden, "AccessDeniedException", err.Error()) + case cerrors.GetCode(err) == cerrors.ResourceExhausted: + // A quota/limit breach maps to Bedrock's ServiceQuotaExceededException. + writeError(w, http.StatusBadRequest, "ServiceQuotaExceededException", err.Error()) case cerrors.IsFailedPrecondition(err): // A resource in a conflicting state (e.g. stopping a terminal job) maps // to Bedrock's ConflictException, matching real AWS. diff --git a/server/aws/bedrockagent/errors.go b/server/aws/bedrockagent/errors.go index b2133e07..ee25fea3 100644 --- a/server/aws/bedrockagent/errors.go +++ b/server/aws/bedrockagent/errors.go @@ -38,6 +38,10 @@ func writeErr(w http.ResponseWriter, err error) { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) case cerrors.IsThrottled(err): writeError(w, http.StatusTooManyRequests, "ThrottlingException", err.Error()) + case cerrors.IsPermissionDenied(err): + writeError(w, http.StatusForbidden, "AccessDeniedException", err.Error()) + case cerrors.GetCode(err) == cerrors.ResourceExhausted: + writeError(w, http.StatusBadRequest, "ServiceQuotaExceededException", err.Error()) default: writeError(w, http.StatusInternalServerError, "InternalServerException", err.Error()) } diff --git a/server/aws/bedrockagentruntime/errors.go b/server/aws/bedrockagentruntime/errors.go index 5a541fc8..31151c1f 100644 --- a/server/aws/bedrockagentruntime/errors.go +++ b/server/aws/bedrockagentruntime/errors.go @@ -38,6 +38,10 @@ func writeErr(w http.ResponseWriter, err error) { writeError(w, http.StatusBadRequest, "ValidationException", err.Error()) case cerrors.IsThrottled(err): writeError(w, http.StatusTooManyRequests, "ThrottlingException", err.Error()) + case cerrors.IsPermissionDenied(err): + writeError(w, http.StatusForbidden, "AccessDeniedException", err.Error()) + case cerrors.GetCode(err) == cerrors.ResourceExhausted: + writeError(w, http.StatusBadRequest, "ServiceQuotaExceededException", err.Error()) default: writeError(w, http.StatusInternalServerError, "InternalServerException", err.Error()) } From 44715f245eb239b31a615f686c1ba4eb7f5cf106 Mon Sep 17 00:00:00 2001 From: Satyam Trivedi Date: Wed, 29 Jul 2026 14:34:20 +0530 Subject: [PATCH 5/5] =?UTF-8?q?fix(bedrock):=20address=20deep=20re-review?= =?UTF-8?q?=20=E2=80=94=20UTF-8=20streaming=20bug,=20list=20determinism,?= =?UTF-8?q?=20aliasing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the deep multi-agent re-review on #214. Correctness: - ConverseStream: chunkText split the completion on a byte offset, corrupting multi-byte UTF-8 responses (broken halves marshalled as U+FFFD). Now advances the split to a rune boundary. Added a multi-byte-UTF-8 streaming roundtrip test (café/日本語/emoji) asserting the reassembled text is valid UTF-8, U+FFFD-free, and equal to non-streaming Converse. Determinism: - ListModelCustomizationJobs, ListCustomModels, ListGuardrails (no-identifier), and ListProvisionedModelThroughputs now use SortedValues() instead of All(), matching the documented memstore contract and the sibling lists. Aliasing / immutability: - FoundationModel Get/List clone the modality/customization/inference slices out (the seed reuses shared backing slices, so a mutated return would corrupt all). - Model-invocation LoggingConfig is deep-copied on store and read (nested S3/CloudWatch pointers no longer shared with caller/store). Fidelity / consistency: - CreateInferenceProfile and CreateAutomatedReasoningPolicy reject duplicate names with AlreadyExists (consistent with prompt routers / provisioned throughput). - UpdateGuardrail rename now errors instead of silently clobbering an existing guardrail with the target name. - RegisterMarketplaceModelEndpoint upserts (registers an externally-created endpoint) instead of requiring the record to pre-exist. - CreateFoundationModelAgreement validates the modelId against the catalog. - CountTokens with neither union member returns ValidationException instead of 0. Docs/tests: documented the bedrock-agent bucket-name-shadowing tradeoff on the handler; added -race tests for the marketplace + AR-policy copy-on-write mutators. Verified: go build/vet/gofmt/golangci-lint clean, go test ./... (204 pkgs) passing, and go test -race on the bedrock/agent packages clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- providers/aws/bedrock/bedrock.go | 23 ++++- providers/aws/bedrock/bedrock_test.go | 39 ++++++++ .../aws/bedrock/concurrency_race_test.go | 96 +++++++++++++++++++ providers/aws/bedrock/management.go | 44 +++++++-- providers/aws/bedrock/management_test.go | 62 ++++++++++++ .../aws/bedrock/marketplace_agreements.go | 30 +++++- .../bedrock/marketplace_agreements_test.go | 46 ++++++++- providers/aws/bedrock/registries.go | 12 +++ providers/aws/bedrock/registries_test.go | 36 +++++++ .../aws/bedrock/counttokens_applyguardrail.go | 6 ++ .../bedrock/sdk_roundtrip_streaming_test.go | 85 ++++++++++++++++ server/aws/bedrock/streaming.go | 13 +++ server/aws/bedrockagent/handler.go | 8 ++ 13 files changed, 485 insertions(+), 15 deletions(-) diff --git a/providers/aws/bedrock/bedrock.go b/providers/aws/bedrock/bedrock.go index 2c09b62b..016a710f 100644 --- a/providers/aws/bedrock/bedrock.go +++ b/providers/aws/bedrock/bedrock.go @@ -80,7 +80,9 @@ func (m *Mock) now() string { // ListFoundationModels returns the seeded foundation-model catalog. func (m *Mock) ListFoundationModels(_ context.Context) ([]driver.FoundationModel, error) { out := make([]driver.FoundationModel, len(m.foundation)) - copy(out, m.foundation) + for i := range m.foundation { + out[i] = cloneFoundationModel(m.foundation[i]) + } return out, nil } @@ -92,11 +94,24 @@ func (m *Mock) GetFoundationModel(_ context.Context, modelID string) (*driver.Fo return nil, errors.Newf(errors.NotFound, "foundation model %q not found", modelID) } - result := *fm + result := cloneFoundationModel(*fm) return &result, nil } +// cloneFoundationModel returns a copy of fm with its slice fields deep-copied so +// callers cannot mutate the shared seed catalog through the returned value. +// +//nolint:gocritic // fm is copied intentionally so slice fields can be reassigned to fresh backing arrays. +func cloneFoundationModel(fm driver.FoundationModel) driver.FoundationModel { + fm.InputModalities = append([]string(nil), fm.InputModalities...) + fm.OutputModalities = append([]string(nil), fm.OutputModalities...) + fm.CustomizationsSupported = append([]string(nil), fm.CustomizationsSupported...) + fm.InferenceTypesSupported = append([]string(nil), fm.InferenceTypesSupported...) + + return fm +} + // findFoundation returns the seeded model matching id by ModelID or ModelARN. func (m *Mock) findFoundation(id string) *driver.FoundationModel { for i := range m.foundation { @@ -198,7 +213,7 @@ func (m *Mock) GetModelCustomizationJob(_ context.Context, jobIdentifier string) // ListModelCustomizationJobs lists all customization jobs. func (m *Mock) ListModelCustomizationJobs(_ context.Context) ([]driver.CustomizationJob, error) { - all := m.jobs.All() + all := m.jobs.SortedValues() out := make([]driver.CustomizationJob, 0, len(all)) for _, job := range all { @@ -212,7 +227,7 @@ func (m *Mock) ListModelCustomizationJobs(_ context.Context) ([]driver.Customiza // ListCustomModels lists all custom models. func (m *Mock) ListCustomModels(_ context.Context) ([]driver.CustomModel, error) { - all := m.models.All() + all := m.models.SortedValues() out := make([]driver.CustomModel, 0, len(all)) for _, cm := range all { diff --git a/providers/aws/bedrock/bedrock_test.go b/providers/aws/bedrock/bedrock_test.go index 500c0226..306f87d0 100644 --- a/providers/aws/bedrock/bedrock_test.go +++ b/providers/aws/bedrock/bedrock_test.go @@ -52,6 +52,45 @@ func TestGetFoundationModel(t *testing.T) { assertError(t, err, true) } +func TestFoundationModelCopyOut(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + // Mutating a returned model's slice must not corrupt the shared seed. + fm, err := m.GetFoundationModel(ctx, titanModel) + requireNoError(t, err) + if len(fm.InputModalities) == 0 { + t.Fatal("expected seeded input modalities") + } + fm.InputModalities[0] = "MUTATED" + + again, err := m.GetFoundationModel(ctx, titanModel) + requireNoError(t, err) + assertEqual(t, "TEXT", again.InputModalities[0]) + + // A different model sharing the same backing seed slice is also unaffected. + other, err := m.GetFoundationModel(ctx, "anthropic.claude-3-sonnet-20240229-v1:0") + requireNoError(t, err) + assertEqual(t, "TEXT", other.InputModalities[0]) + + // ListFoundationModels returns independent copies too. + list, err := m.ListFoundationModels(ctx) + requireNoError(t, err) + for i := range list { + if len(list[i].InputModalities) > 0 { + list[i].InputModalities[0] = "MUTATED" + } + } + + relist, err := m.ListFoundationModels(ctx) + requireNoError(t, err) + for i := range relist { + if len(relist[i].InputModalities) > 0 { + assertEqual(t, "TEXT", relist[i].InputModalities[0]) + } + } +} + func TestCreateModelCustomizationJob(t *testing.T) { tests := []struct { name string diff --git a/providers/aws/bedrock/concurrency_race_test.go b/providers/aws/bedrock/concurrency_race_test.go index 32be3261..a523ca54 100644 --- a/providers/aws/bedrock/concurrency_race_test.go +++ b/providers/aws/bedrock/concurrency_race_test.go @@ -107,3 +107,99 @@ func TestConcurrentEvalJobAccessRace(t *testing.T) { wg.Wait() } + +// TestConcurrentMarketplaceEndpointAccessRace hammers the marketplace endpoint +// copy-on-write mutators (Register/Update) concurrently with Get/List reads, +// exercising the copy-on-write paths under `go test -race`. +func TestConcurrentMarketplaceEndpointAccessRace(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + created, err := m.CreateMarketplaceModelEndpoint(ctx, bedrockdriver.MarketplaceEndpointConfig{ + EndpointName: "race-endpoint", + ModelSourceIdentifier: "arn:aws:sagemaker:us-east-1:aws:hub-content/model/1", + EndpointConfig: []byte(`{"sageMaker":{"instanceType":"ml.m5.large"}}`), + }) + if err != nil { + t.Fatalf("CreateMarketplaceModelEndpoint: %v", err) + } + + arn := created.EndpointARN + + const workers = 8 + + const iters = 50 + + var wg sync.WaitGroup + + for w := 0; w < workers; w++ { + wg.Add(1) + + go func(id int) { + defer wg.Done() + + for i := 0; i < iters; i++ { + switch id % 4 { + case 0: + _, _ = m.RegisterMarketplaceModelEndpoint(ctx, arn, "arn:model/source-race") + case 1: + _, _ = m.UpdateMarketplaceModelEndpoint(ctx, arn, []byte(`{"sageMaker":{"instanceType":"ml.m5.xlarge"}}`)) + case 2: + _, _ = m.GetMarketplaceModelEndpoint(ctx, arn) + default: + _, _ = m.ListMarketplaceModelEndpoints(ctx) + } + } + }(w) + } + + wg.Wait() +} + +// TestConcurrentARPolicyAccessRace hammers the automated-reasoning-policy +// copy-on-write Update mutator concurrently with Get/List reads, exercising the +// copy-on-write path under `go test -race`. +func TestConcurrentARPolicyAccessRace(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + policy, err := m.CreateAutomatedReasoningPolicy(ctx, bedrockdriver.AutomatedReasoningPolicyConfig{ + Name: "race-policy", + PolicyDefinition: []byte(`{"rules":[]}`), + }) + if err != nil { + t.Fatalf("CreateAutomatedReasoningPolicy: %v", err) + } + + arn := policy.ARN + + const workers = 8 + + const iters = 50 + + var wg sync.WaitGroup + + for w := 0; w < workers; w++ { + wg.Add(1) + + go func(id int) { + defer wg.Done() + + for i := 0; i < iters; i++ { + switch id % 3 { + case 0: + _, _ = m.UpdateAutomatedReasoningPolicy(ctx, arn, bedrockdriver.AutomatedReasoningPolicyUpdate{ + Description: "updated", + PolicyDefinition: []byte(`{"rules":[{"id":"r1"}]}`), + }) + case 1: + _, _ = m.GetAutomatedReasoningPolicy(ctx, arn) + default: + _, _ = m.ListAutomatedReasoningPolicies(ctx) + } + } + }(w) + } + + wg.Wait() +} diff --git a/providers/aws/bedrock/management.go b/providers/aws/bedrock/management.go index fcee0019..38722697 100644 --- a/providers/aws/bedrock/management.go +++ b/providers/aws/bedrock/management.go @@ -89,7 +89,7 @@ func (m *Mock) ListGuardrails(_ context.Context, identifier string) ([]driver.Gu return rec.allSnapshots(), nil } - all := m.guardrails.All() + all := m.guardrails.SortedValues() out := make([]driver.Guardrail, 0, len(all)) for _, rec := range all { @@ -112,13 +112,20 @@ func (m *Mock) UpdateGuardrail(_ context.Context, identifier string, cfg driver. rec.mu.Lock() g := rec.draft oldName := g.Name - g.Name = orDefault(cfg.Name, g.Name) + + newName := orDefault(cfg.Name, oldName) + if newName != oldName && m.guardrails.Has(newName) { + rec.mu.Unlock() + + return nil, errors.Newf(errors.AlreadyExists, "guardrail %q already exists", newName) + } + + g.Name = newName g.Description = cfg.Description g.BlockedInputMessaging = orDefault(cfg.BlockedInputMessaging, g.BlockedInputMessaging) g.BlockedOutputsMessaging = orDefault(cfg.BlockedOutputsMessaging, g.BlockedOutputsMessaging) g.GuardrailPolicies = deepCopyGuardrailPolicies(cfg.GuardrailPolicies) g.UpdatedAt = m.now() - newName := g.Name result := cloneGuardrail(g) rec.mu.Unlock() @@ -255,7 +262,7 @@ func (m *Mock) GetProvisionedModelThroughput(_ context.Context, identifier strin // ListProvisionedModelThroughputs lists all provisioned throughputs. func (m *Mock) ListProvisionedModelThroughputs(_ context.Context) ([]driver.ProvisionedThroughput, error) { - all := m.provisioned.All() + all := m.provisioned.SortedValues() out := make([]driver.ProvisionedThroughput, 0, len(all)) for _, pt := range all { @@ -309,7 +316,7 @@ func (m *Mock) resolveModelARN(id string) string { // PutModelInvocationLoggingConfiguration sets the invocation logging config. func (m *Mock) PutModelInvocationLoggingConfiguration(_ context.Context, cfg driver.LoggingConfig) error { - stored := cfg + stored := deepCopyLoggingConfig(cfg) m.logMu.Lock() m.logging = &stored @@ -318,6 +325,31 @@ func (m *Mock) PutModelInvocationLoggingConfiguration(_ context.Context, cfg dri return nil } +// deepCopyLoggingConfig returns a copy of cfg with its nested pointer fields +// freshly allocated, so stored and returned configs never alias a caller's +// pointers. Nil pointers stay nil. +func deepCopyLoggingConfig(cfg driver.LoggingConfig) driver.LoggingConfig { + out := cfg + + if cfg.S3 != nil { + s3 := *cfg.S3 + out.S3 = &s3 + } + + if cfg.CloudWatch != nil { + cw := *cfg.CloudWatch + + if cfg.CloudWatch.LargeDataDeliveryS3 != nil { + ldd := *cfg.CloudWatch.LargeDataDeliveryS3 + cw.LargeDataDeliveryS3 = &ldd + } + + out.CloudWatch = &cw + } + + return out +} + // GetModelInvocationLoggingConfiguration returns the invocation logging config, // or nil if none is set. func (m *Mock) GetModelInvocationLoggingConfiguration(_ context.Context) (*driver.LoggingConfig, error) { @@ -328,7 +360,7 @@ func (m *Mock) GetModelInvocationLoggingConfiguration(_ context.Context) (*drive return nil, nil //nolint:nilnil // an unset logging config is a valid, non-error result } - result := *m.logging + result := deepCopyLoggingConfig(*m.logging) return &result, nil } diff --git a/providers/aws/bedrock/management_test.go b/providers/aws/bedrock/management_test.go index f5657d99..4cf63aba 100644 --- a/providers/aws/bedrock/management_test.go +++ b/providers/aws/bedrock/management_test.go @@ -4,6 +4,7 @@ import ( "context" "testing" + "github.com/stackshy/cloudemu/v2/errors" bedrockdriver "github.com/stackshy/cloudemu/v2/services/bedrock/driver" ) @@ -81,6 +82,39 @@ func TestGuardrailRenameRekeys(t *testing.T) { assertEqual(t, 1, len(list)) } +func TestGuardrailRenameCollision(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + a := newGuardrail(t, m, "gr-a") + newGuardrail(t, m, "gr-b") + + // Renaming A onto B's name must fail with AlreadyExists and leave both intact. + _, err := m.UpdateGuardrail(ctx, a.ID, bedrockdriver.GuardrailConfig{ + Name: "gr-b", + BlockedInputMessaging: "in", + BlockedOutputsMessaging: "out", + }) + assertError(t, err, true) + if !errors.IsAlreadyExists(err) { + t.Fatalf("expected AlreadyExists, got %v", err) + } + + // A still exists under its original name, unmodified. + gotA, err := m.GetGuardrail(ctx, a.ID, "") + requireNoError(t, err) + assertEqual(t, "gr-a", gotA.Name) + + // B is intact. + gotB, err := m.GetGuardrail(ctx, "gr-b", "") + requireNoError(t, err) + assertEqual(t, "gr-b", gotB.Name) + + list, err := m.ListGuardrails(ctx, "") + requireNoError(t, err) + assertEqual(t, 2, len(list)) +} + func TestGuardrailValidation(t *testing.T) { m := newTestMock() ctx := context.Background() @@ -277,3 +311,31 @@ func TestModelInvocationLogging(t *testing.T) { t.Fatalf("expected nil after delete, got %+v", cfg) } } + +func TestModelInvocationLoggingCopyOut(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + // The stored config must not alias the caller's nested pointers. + in := bedrockdriver.LoggingConfig{ + S3: &bedrockdriver.S3LoggingConfig{BucketName: "logs", KeyPrefix: "p/"}, + } + requireNoError(t, m.PutModelInvocationLoggingConfiguration(ctx, in)) + + // Mutating the caller's config after Put must not change stored state. + in.S3.BucketName = "caller-mutated" + + cfg, err := m.GetModelInvocationLoggingConfiguration(ctx) + requireNoError(t, err) + if cfg == nil || cfg.S3 == nil { + t.Fatal("expected logging config with S3") + } + assertEqual(t, "logs", cfg.S3.BucketName) + + // Mutating a returned config must not corrupt state for the next reader. + cfg.S3.BucketName = "reader-mutated" + + again, err := m.GetModelInvocationLoggingConfiguration(ctx) + requireNoError(t, err) + assertEqual(t, "logs", again.S3.BucketName) +} diff --git a/providers/aws/bedrock/marketplace_agreements.go b/providers/aws/bedrock/marketplace_agreements.go index 93b4c0b9..ac7f7903 100644 --- a/providers/aws/bedrock/marketplace_agreements.go +++ b/providers/aws/bedrock/marketplace_agreements.go @@ -110,8 +110,10 @@ func (m *Mock) DeleteMarketplaceModelEndpoint(_ context.Context, endpointARN str return nil } -// RegisterMarketplaceModelEndpoint registers an existing endpoint, marking it -// REGISTERED. +// RegisterMarketplaceModelEndpoint registers an externally-created SageMaker +// endpoint with Bedrock, marking it REGISTERED. This is an upsert: if the +// endpoint is not already tracked it is created as a new REGISTERED record keyed +// by endpointIdentifier; if it exists, its model source and status are updated. func (m *Mock) RegisterMarketplaceModelEndpoint( _ context.Context, endpointIdentifier, modelSourceIdentifier string, ) (*driver.MarketplaceEndpoint, error) { @@ -119,16 +121,32 @@ func (m *Mock) RegisterMarketplaceModelEndpoint( return nil, errors.New(errors.InvalidArgument, "modelSourceIdentifier is required") } + now := m.now() + stored, ok := m.marketplaceEndpoints.Get(endpointIdentifier) if !ok { - return nil, errors.Newf(errors.NotFound, "marketplace model endpoint %q not found", endpointIdentifier) + // Not tracked yet: register the externally-created endpoint as a new + // REGISTERED record keyed by its identifier. + endpoint := &driver.MarketplaceEndpoint{ + EndpointARN: endpointIdentifier, + ModelSourceIdentifier: modelSourceIdentifier, + EndpointStatus: driver.MarketplaceEndpointStatusInService, + Status: driver.MarketplaceEndpointStatusRegistered, + CreatedAt: now, + UpdatedAt: now, + } + m.marketplaceEndpoints.Set(endpointIdentifier, endpoint) + + result := cloneMarketplaceEndpoint(endpoint) + + return &result, nil } // Copy-on-write: mutate a copy so concurrent readers never race the write. updated := *stored updated.ModelSourceIdentifier = modelSourceIdentifier updated.Status = driver.MarketplaceEndpointStatusRegistered - updated.UpdatedAt = m.now() + updated.UpdatedAt = now m.marketplaceEndpoints.Set(endpointIdentifier, &updated) result := cloneMarketplaceEndpoint(&updated) @@ -161,6 +179,10 @@ func (m *Mock) CreateFoundationModelAgreement(_ context.Context, modelID, offerT return "", errors.New(errors.InvalidArgument, "offerToken is required") } + if m.findFoundation(modelID) == nil { + return "", errors.Newf(errors.InvalidArgument, "foundation model %q not found", modelID) + } + m.fmAgreements.Set(modelID, true) return modelID, nil diff --git a/providers/aws/bedrock/marketplace_agreements_test.go b/providers/aws/bedrock/marketplace_agreements_test.go index b84cab2f..76e7bd56 100644 --- a/providers/aws/bedrock/marketplace_agreements_test.go +++ b/providers/aws/bedrock/marketplace_agreements_test.go @@ -90,13 +90,41 @@ func TestMarketplaceEndpointValidationAndErrors(t *testing.T) { _, err = m.UpdateMarketplaceModelEndpoint(ctx, "arn:missing", []byte(`{}`)) assertError(t, err, true) - _, err = m.RegisterMarketplaceModelEndpoint(ctx, "arn:missing", "arn:model/s") + // Register requires a model source identifier. + _, err = m.RegisterMarketplaceModelEndpoint(ctx, "arn:some-endpoint", "") assertError(t, err, true) assertError(t, m.DeregisterMarketplaceModelEndpoint(ctx, "arn:missing"), true) assertError(t, m.DeleteMarketplaceModelEndpoint(ctx, "arn:missing"), true) } +// TestRegisterMarketplaceEndpointUpsert verifies that registering a brand-new +// (untracked) endpoint identifier creates a REGISTERED record retrievable by Get. +func TestRegisterMarketplaceEndpointUpsert(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + const arn = "arn:aws:sagemaker:us-east-1:123456789012:endpoint/external-endpoint" + + reg, err := m.RegisterMarketplaceModelEndpoint(ctx, arn, "arn:model/source-1") + requireNoError(t, err) + assertEqual(t, arn, reg.EndpointARN) + assertEqual(t, "arn:model/source-1", reg.ModelSourceIdentifier) + assertEqual(t, bedrockdriver.MarketplaceEndpointStatusRegistered, reg.Status) + assertEqual(t, bedrockdriver.MarketplaceEndpointStatusInService, reg.EndpointStatus) + + got, err := m.GetMarketplaceModelEndpoint(ctx, arn) + requireNoError(t, err) + assertEqual(t, arn, got.EndpointARN) + assertEqual(t, "arn:model/source-1", got.ModelSourceIdentifier) + assertEqual(t, bedrockdriver.MarketplaceEndpointStatusRegistered, got.Status) + + // Re-registering the now-tracked endpoint updates its model source. + reg2, err := m.RegisterMarketplaceModelEndpoint(ctx, arn, "arn:model/source-2") + requireNoError(t, err) + assertEqual(t, "arn:model/source-2", reg2.ModelSourceIdentifier) +} + // TestMarketplaceEndpointCopyOut verifies that mutating the EndpointConfig // bytes of a returned endpoint does not affect the stored value. func TestMarketplaceEndpointCopyOut(t *testing.T) { @@ -166,6 +194,22 @@ func TestFoundationModelAgreementValidation(t *testing.T) { _, err = m.CreateFoundationModelAgreement(ctx, titanModel, "") assertError(t, err, true) + // A model that is not in the catalog is rejected. + _, err = m.CreateFoundationModelAgreement(ctx, "bogus.nonexistent-model-v1", "token") + assertError(t, err, true) + + // A real catalog model (fetched from ListFoundationModels) succeeds. + models, err := m.ListFoundationModels(ctx) + requireNoError(t, err) + + if len(models) == 0 { + t.Fatal("expected a non-empty foundation model catalog") + } + + modelID, err := m.CreateFoundationModelAgreement(ctx, models[0].ModelID, "token") + requireNoError(t, err) + assertEqual(t, models[0].ModelID, modelID) + _, err = m.ListFoundationModelAgreementOffers(ctx, "", "ALL") assertError(t, err, true) diff --git a/providers/aws/bedrock/registries.go b/providers/aws/bedrock/registries.go index 88e5f643..43e02acf 100644 --- a/providers/aws/bedrock/registries.go +++ b/providers/aws/bedrock/registries.go @@ -22,6 +22,12 @@ func (m *Mock) CreateInferenceProfile(_ context.Context, cfg driver.InferencePro return nil, errors.New(errors.InvalidArgument, "modelSource.copyFrom is required") } + for _, existing := range m.inferenceProfiles.SortedValues() { + if existing.Name == cfg.Name { + return nil, errors.Newf(errors.AlreadyExists, "inference profile %q already exists", cfg.Name) + } + } + now := m.now() id := idgen.GenerateID("") arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "application-inference-profile/"+id) @@ -191,6 +197,12 @@ func (m *Mock) CreateAutomatedReasoningPolicy( return nil, errors.New(errors.InvalidArgument, "name is required") } + for _, existing := range m.arPolicies.SortedValues() { + if existing.Name == cfg.Name { + return nil, errors.Newf(errors.AlreadyExists, "automated reasoning policy %q already exists", cfg.Name) + } + } + now := m.now() id := idgen.GenerateID("") arn := idgen.AWSARN("bedrock", m.opts.Region, m.opts.AccountID, "automated-reasoning-policy/"+id) diff --git a/providers/aws/bedrock/registries_test.go b/providers/aws/bedrock/registries_test.go index 9ecb1095..7e30bedf 100644 --- a/providers/aws/bedrock/registries_test.go +++ b/providers/aws/bedrock/registries_test.go @@ -61,6 +61,24 @@ func TestInferenceProfileValidation(t *testing.T) { assertError(t, m.DeleteInferenceProfile(ctx, "missing"), true) } +// TestInferenceProfileDuplicateName verifies a second create with the same name +// returns AlreadyExists. +func TestInferenceProfileDuplicateName(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + cfg := bedrockdriver.InferenceProfileConfig{ + Name: "dup-profile", + ModelSourceCopyFrom: "arn:aws:bedrock:us-east-1:123456789012:foundation-model/" + titanModel, + } + + _, err := m.CreateInferenceProfile(ctx, cfg) + requireNoError(t, err) + + _, err = m.CreateInferenceProfile(ctx, cfg) + assertError(t, err, true) +} + func TestPromptRouterLifecycle(t *testing.T) { m := newTestMock() ctx := context.Background() @@ -179,3 +197,21 @@ func TestAutomatedReasoningPolicyValidationAndErrors(t *testing.T) { assertError(t, m.DeleteAutomatedReasoningPolicy(ctx, "arn:missing"), true) } + +// TestAutomatedReasoningPolicyDuplicateName verifies a second create with the +// same name returns AlreadyExists. +func TestAutomatedReasoningPolicyDuplicateName(t *testing.T) { + m := newTestMock() + ctx := context.Background() + + cfg := bedrockdriver.AutomatedReasoningPolicyConfig{ + Name: "dup-policy", + PolicyDefinition: []byte(`{"rules":[]}`), + } + + _, err := m.CreateAutomatedReasoningPolicy(ctx, cfg) + requireNoError(t, err) + + _, err = m.CreateAutomatedReasoningPolicy(ctx, cfg) + assertError(t, err, true) +} diff --git a/server/aws/bedrock/counttokens_applyguardrail.go b/server/aws/bedrock/counttokens_applyguardrail.go index 41eab21b..d0e2dd6d 100644 --- a/server/aws/bedrock/counttokens_applyguardrail.go +++ b/server/aws/bedrock/counttokens_applyguardrail.go @@ -21,6 +21,12 @@ func (h *Handler) countTokens(w http.ResponseWriter, r *http.Request, modelID st } req := body.Input + if req.InvokeModel == nil && req.Converse == nil { + writeError(w, http.StatusBadRequest, "ValidationException", "input must specify either converse or invokeModel") + + return + } + in := bedrockdriver.CountTokensInput{ModelID: modelID} switch { diff --git a/server/aws/bedrock/sdk_roundtrip_streaming_test.go b/server/aws/bedrock/sdk_roundtrip_streaming_test.go index fe5e5bcd..98e8dfcf 100644 --- a/server/aws/bedrock/sdk_roundtrip_streaming_test.go +++ b/server/aws/bedrock/sdk_roundtrip_streaming_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "strings" "testing" + "unicode/utf8" "github.com/aws/aws-sdk-go-v2/aws" awsruntime "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" @@ -73,6 +74,90 @@ func TestSDKConverseStream(t *testing.T) { } } +// TestSDKConverseStreamMultibyteRuneBoundary streams a completion built from a +// prompt dense with multi-byte UTF-8 runes and long enough that the delta split +// point falls in the middle of a multi-byte rune. It asserts the reassembled +// streamed text is valid UTF-8, carries no U+FFFD replacement characters, and +// contains the multi-byte prompt substring byte-for-byte. It also compares the +// streamed text against the non-streaming Converse text for the same prompt. +func TestSDKConverseStreamMultibyteRuneBoundary(t *testing.T) { + client := newRuntimeClient(t) + + // Repeated five times so the byte midpoint of the echoed completion lands + // mid-rune, exercising the rune-boundary split (a raw byte split here would + // corrupt both halves into U+FFFD). + prompt := strings.Repeat("café ☕ 日本語 🎉 naïve résumé ", 5) + + msg := []runtimetypes.Message{ + { + Role: runtimetypes.ConversationRoleUser, + Content: []runtimetypes.ContentBlock{&runtimetypes.ContentBlockMemberText{Value: prompt}}, + }, + } + + streamOut, err := client.ConverseStream(context.Background(), &awsruntime.ConverseStreamInput{ + ModelId: aws.String(claudeModel), + Messages: msg, + }) + if err != nil { + t.Fatalf("ConverseStream: %v", err) + } + + stream := streamOut.GetStream() + defer stream.Close() + + var streamed strings.Builder + + for ev := range stream.Events() { + if d, ok := ev.(*runtimetypes.ConverseStreamOutputMemberContentBlockDelta); ok { + if td, ok := d.Value.Delta.(*runtimetypes.ContentBlockDeltaMemberText); ok { + streamed.WriteString(td.Value) + } + } + } + + if err := stream.Err(); err != nil { + t.Fatalf("stream error: %v", err) + } + + got := streamed.String() + + if !utf8.ValidString(got) { + t.Fatalf("streamed text is not valid UTF-8: %q", got) + } + + if strings.ContainsRune(got, '�') { + t.Fatalf("streamed text contains U+FFFD replacement character: %q", got) + } + + if !strings.Contains(got, prompt) { + t.Fatalf("streamed text does not contain the multi-byte prompt substring\n got: %q\nwant substring: %q", got, prompt) + } + + // The non-streaming Converse of the same prompt must yield identical text. + convOut, err := client.Converse(context.Background(), &awsruntime.ConverseInput{ + ModelId: aws.String(claudeModel), + Messages: msg, + }) + if err != nil { + t.Fatalf("Converse: %v", err) + } + + var nonStreamed strings.Builder + + if out, ok := convOut.Output.(*runtimetypes.ConverseOutputMemberMessage); ok { + for _, block := range out.Value.Content { + if tb, ok := block.(*runtimetypes.ContentBlockMemberText); ok { + nonStreamed.WriteString(tb.Value) + } + } + } + + if got != nonStreamed.String() { + t.Fatalf("streamed text != non-streamed text\nstreamed: %q\nnon-streamed: %q", got, nonStreamed.String()) + } +} + func TestSDKInvokeModelWithResponseStream(t *testing.T) { client := newRuntimeClient(t) diff --git a/server/aws/bedrock/streaming.go b/server/aws/bedrock/streaming.go index d5758dab..801f1180 100644 --- a/server/aws/bedrock/streaming.go +++ b/server/aws/bedrock/streaming.go @@ -6,6 +6,7 @@ import ( "io" "net/http" "strings" + "unicode/utf8" "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream" "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi" @@ -135,12 +136,24 @@ func metadataPayload(out *bedrockdriver.ConverseOutput) []byte { // chunkText splits s into up to two contentBlockDelta chunks so the emulated // stream delivers more than one delta when the text is long enough. It always // returns at least one chunk. +// +// The split point is advanced to a UTF-8 rune boundary so neither half can +// contain a truncated multi-byte rune. Splitting on a raw byte offset would +// leave both halves as invalid UTF-8, which encoding/json then marshals as the +// U+FFFD replacement character, corrupting the streamed text. func chunkText(s string) []string { if len(s) < minSplitLen { return []string{s} } mid := len(s) / 2 + for mid < len(s) && !utf8.RuneStart(s[mid]) { + mid++ + } + + if mid == 0 || mid >= len(s) { + return []string{s} + } return []string{s[:mid], s[mid:]} } diff --git a/server/aws/bedrockagent/handler.go b/server/aws/bedrockagent/handler.go index 53d29adf..c4a6916b 100644 --- a/server/aws/bedrockagent/handler.go +++ b/server/aws/bedrockagent/handler.go @@ -78,6 +78,14 @@ func New(drv badriver.BedrockAgent) *Handler { } // Matches claims the Bedrock Agent authoring URL prefixes. +// +// The claimed REST roots /agents, /knowledgebases, /flows, and /prompts are +// anchored via underPrefix. As with the sibling /custom-models and EKS +// /clusters handlers, this means an S3 bucket named EXACTLY agents, +// knowledgebases, flows, or prompts under path-style addressing is claimed +// here before reaching the S3 catch-all. This is the accepted REST-vs-catch-all +// tradeoff; callers needing those exact bucket names should use +// virtual-host-style addressing. func (*Handler) Matches(r *http.Request) bool { p := r.URL.Path