From e05bfce39fbaa7905c79f891a09903d6a4e902a7 Mon Sep 17 00:00:00 2001 From: bgagent Date: Mon, 27 Jul 2026 13:17:42 -0400 Subject: [PATCH] feat(registry): resolve + load registry assets into tasks (#246) Builds on the catalog PR to actually consume registry assets at task time: - Orchestrator resolve-step (`resolveRegistryAssets`): resolves a blueprint's `registry://` mcp_server / cedar_policy_module / skill refs at task start, fail-closed; stamps the `{kind,id,version}` triples on the TaskRecord for audit, merges resolved cedar_text into `cedar_policies`, and threads the runtime bundle into the agent payload. - Blueprint asset props + onUpdate fix: `assets.{mcpServers,cedarPolicyModules, skills}` with `RegistryRefValidation`; the three onUpdate helpers now write the asset-ref columns so redeploying an onboarded repo no longer drops them. - Agent loaders (registry.loader): mcp_server merges into `.mcp.json`; cedar_policy_module flows through PolicyEngine's unannotated `extra_policies`; skill prompt fragments append to the system prompt (build_skill_prompt_fragment). - TaskOrchestrator IAM: read-only bedrock-agentcore registry access so the orchestrator can resolve refs. Depends on the catalog PR (feat/246-registry-catalog): imports the RegistryClient port, ref grammar, and resolver from that branch. --- agent/src/models.py | 7 +- agent/src/pipeline.py | 13 ++ agent/src/prompt_builder.py | 12 ++ agent/src/registry/loader.py | 148 +++++++++++++++++ agent/src/server.py | 6 + agent/tests/test_registry_loader.py | 151 ++++++++++++++++++ cdk/src/constructs/blueprint.ts | 73 +++++++++ cdk/src/constructs/task-orchestrator.ts | 38 ++++- cdk/src/handlers/shared/orchestrator.ts | 81 +++++++++- cdk/src/handlers/shared/repo-config.ts | 21 +++ cdk/src/stacks/agent.ts | 18 +++ cdk/test/constructs/blueprint.test.ts | 67 ++++++++ .../shared/registry-orchestrator.test.ts | 141 ++++++++++++++++ 13 files changed, 772 insertions(+), 4 deletions(-) create mode 100644 agent/src/registry/loader.py create mode 100644 agent/tests/test_registry_loader.py create mode 100644 cdk/test/handlers/shared/registry-orchestrator.test.ts diff --git a/agent/src/models.py b/agent/src/models.py index d7687c4e7..befbd14ea 100644 --- a/agent/src/models.py +++ b/agent/src/models.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Literal, Self +from typing import Any, Literal, Self from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -221,6 +221,11 @@ class TaskConfig(BaseModel): trace: bool = False # Enriched mid-flight by pipeline.py: cedar_policies: list[str] = [] + # Registry assets (#246) resolved by the orchestrator and threaded in the + # payload. Each entry is ``{kind, namespace, name, version, runtime}``; the + # per-kind loaders (registry.loader) apply them — mcp_server merges into + # ``.mcp.json`` (PR 2); cedar_policy_module / skill land in PR 3. + resolved_assets: list[dict[str, Any]] = Field(default_factory=list) # Cedar HITL (§7.3, §10.2). Per-task approval defaults threaded # from the orchestrator payload; consumed by PolicyEngine at # construction so the engine seeds ApprovalAllowlist and adopts diff --git a/agent/src/pipeline.py b/agent/src/pipeline.py index 78b2972ac..31da593d7 100644 --- a/agent/src/pipeline.py +++ b/agent/src/pipeline.py @@ -626,6 +626,7 @@ def run_task( trace: bool = False, user_id: str = "", attachments: list[dict] | None = None, + resolved_assets: list[dict] | None = None, ) -> dict: """Run the full agent pipeline and return a serialized result dict. @@ -676,6 +677,11 @@ def run_task( if cedar_policies: config.cedar_policies = cedar_policies + # Registry assets (#246) resolved by the orchestrator — applied by the + # per-kind loaders below (mcp_server → .mcp.json in PR 2). + if resolved_assets: + config.resolved_assets = resolved_assets + # Export session-tag values so tenant-data boto3 clients (DDB/S3) assume # the per-task SessionRole with {user_id, repo, task_id} tags. No-op when # AGENT_SESSION_ROLE_ARN is unset (local/dev/tests). @@ -915,6 +921,13 @@ def _on_trace_truncated(max_bytes: int, first_dropped: int) -> None: # the user gets immediate feedback; see the Early ACK block above.) configure_channel_mcp(setup.repo_dir, config.channel_source) + # Registry assets (#246): merge resolved mcp_server configs into + # .mcp.json alongside the channel MCP entry, before the project scan. + if config.resolved_assets: + from registry.loader import apply_resolved_assets + + apply_resolved_assets(setup.repo_dir, config.resolved_assets) + # Download attachments from S3 (version-pinned, integrity-verified) prepared_attachments: list = [] if config.attachments: diff --git a/agent/src/prompt_builder.py b/agent/src/prompt_builder.py index 6a262cbba..3714c171e 100644 --- a/agent/src/prompt_builder.py +++ b/agent/src/prompt_builder.py @@ -76,6 +76,13 @@ def build_system_prompt( if channel_addendum: system_prompt += channel_addendum + # Registry skill assets (#246, PR 3): append resolved prompt fragments. Placed + # after channel guidance so operator-attached skills sit at the recency end. + if config.resolved_assets: + from registry.loader import build_skill_prompt_fragment + + system_prompt += build_skill_prompt_fragment(config.resolved_assets) + return system_prompt @@ -109,6 +116,11 @@ def build_repoless_system_prompt( if channel_addendum: system_prompt += channel_addendum + if config.resolved_assets: + from registry.loader import build_skill_prompt_fragment + + system_prompt += build_skill_prompt_fragment(config.resolved_assets) + return system_prompt diff --git a/agent/src/registry/loader.py b/agent/src/registry/loader.py new file mode 100644 index 000000000..284b011ab --- /dev/null +++ b/agent/src/registry/loader.py @@ -0,0 +1,148 @@ +"""Apply resolved registry assets (#246) to the agent's runtime environment. + +The orchestrator resolves the Blueprint's ``registry://`` refs and threads a +bundle of ``{kind, namespace, name, version, runtime}`` entries into the payload +(``TaskConfig.resolved_assets``). Each per-kind loader here applies its runtime +payload: + + * ``mcp_server`` → merge the connection config into ``.mcp.json`` (PR 2). + * ``cedar_policy_module`` / ``skill`` → PR 3. + +The merge mirrors ``channel_mcp.configure_channel_mcp``: read the existing +``.mcp.json`` (if any), overlay the registry servers without clobbering other +entries, and write it back. Runs alongside the channel MCP wiring so the SDK's +project-scoped scan picks up both. +""" + +from __future__ import annotations + +import json +import os +from typing import Any + +from shell import log + +# The runtime payload for an mcp_server asset is a single ``mcpServers`` entry's +# value (transport/url/headers/…); we key it by ``__`` so two +# registry servers never collide and the source asset is legible in the config. +_MCP_KIND = "mcp_server" +_SKILL_KIND = "skill" + + +def _server_key(asset: dict[str, Any]) -> str: + namespace = asset.get("namespace", "") + name = asset.get("name", "") + return f"{namespace}__{name}".replace("-", "_") + + +def _read_existing_mcp_config(path: str) -> dict[str, Any]: + """Return the parsed .mcp.json at ``path``, or {} if absent/invalid. + + Mirrors ``channel_mcp._read_existing_mcp_config`` — a malformed file is + logged and treated as absent rather than crashing the agent. + """ + if not os.path.isfile(path): + return {} + try: + with open(path, encoding="utf-8") as f: + parsed = json.load(f) + if isinstance(parsed, dict): + return parsed + log("WARN", f"Ignoring non-object .mcp.json at {path} (got {type(parsed).__name__})") + except (OSError, json.JSONDecodeError) as e: + log("WARN", f"Failed to read existing .mcp.json at {path}: {type(e).__name__}: {e}") + return {} + + +def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> int: + """Merge resolved ``mcp_server`` assets into ``/.mcp.json``. + + Returns the number of MCP servers written. A no-op (returns 0) when there are + no mcp_server assets or ``repo_dir`` is missing. + """ + mcp_assets = [a for a in resolved_assets if a.get("kind") == _MCP_KIND] + if not mcp_assets: + return 0 + + if not repo_dir or not os.path.isdir(repo_dir): + log("WARN", f"apply_mcp_assets: repo_dir missing or not a directory: {repo_dir!r}") + return 0 + + mcp_path = os.path.join(repo_dir, ".mcp.json") + config = _read_existing_mcp_config(mcp_path) + servers = config.get("mcpServers") + if not isinstance(servers, dict): + servers = {} + + written = 0 + for asset in mcp_assets: + runtime = asset.get("runtime") + if not isinstance(runtime, dict) or not runtime: + log("WARN", f"apply_mcp_assets: skipping {_server_key(asset)} — empty runtime payload") + continue + servers[_server_key(asset)] = runtime + written += 1 + + if written == 0: + return 0 + + config["mcpServers"] = servers + try: + with open(mcp_path, "w", encoding="utf-8") as f: + json.dump(config, f, indent=2) + f.write("\n") + except OSError as e: + log("ERROR", f"apply_mcp_assets: failed to write {mcp_path}: {e}") + return 0 + + log("TASK", f"Registry: merged {written} MCP server(s) into {mcp_path}") + return written + + +def build_skill_prompt_fragment(resolved_assets: list[dict[str, Any]]) -> str: + """Assemble the appended system-prompt text from resolved ``skill`` assets. + + Each skill's runtime payload carries a ``prompt_fragment`` (and optional + advisory ``tool_hints``). Fragments are concatenated in resolution order under + a single heading, so the model sees them as extra instructions. Returns "" when + there are no skills — the caller then appends nothing. + + Skills are prompt text only: a skill cannot invoke tools; ``tool_hints`` are + advisory prose referencing tools an MCP server separately provides (no + transitive dependency — the operator attaches both). + """ + skills = [a for a in resolved_assets if a.get("kind") == _SKILL_KIND] + if not skills: + return "" + + parts: list[str] = [] + for asset in skills: + runtime = asset.get("runtime") + if not isinstance(runtime, dict): + continue + fragment = runtime.get("prompt_fragment") + if not isinstance(fragment, str) or not fragment.strip(): + continue + name = f"{asset.get('namespace', '')}/{asset.get('name', '')}" + parts.append(f"### Skill: {name}\n\n{fragment.strip()}") + hints = runtime.get("tool_hints") + if isinstance(hints, list) and hints: + parts.append(f"_Suggested tools: {', '.join(str(h) for h in hints)}._") + + if not parts: + return "" + body = "\n\n".join(parts) + log("TASK", f"Registry: appended {len(skills)} skill fragment(s) to the system prompt") + return f"\n\n## Skills\n\n{body}" + + +def apply_resolved_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> None: + """Apply the asset kinds that mutate on-disk state (mcp_server → .mcp.json). + + Cedar policy modules are applied orchestrator-side (merged into the + cedar_policies payload) and skills are applied in prompt_builder via + :func:`build_skill_prompt_fragment`, so neither is handled here. + """ + if not resolved_assets: + return + apply_mcp_assets(repo_dir, resolved_assets) diff --git a/agent/src/server.py b/agent/src/server.py index 9045716a4..de83dc209 100644 --- a/agent/src/server.py +++ b/agent/src/server.py @@ -406,6 +406,7 @@ def _run_task_background( user_id: str = "", workload_access_token: str = "", attachments: list[dict] | None = None, + resolved_assets: list[dict] | None = None, ) -> None: """Run the agent task in a background thread.""" global _background_pipeline_failed @@ -491,6 +492,7 @@ def _run_task_background( trace=trace, user_id=user_id, attachments=attachments, + resolved_assets=resolved_assets, ) _background_pipeline_failed = False except Exception as e: @@ -536,6 +538,9 @@ def _extract_invocation_params(inp: dict, request: Request) -> dict: branch_name = inp.get("branch_name", "") pr_number = str(inp.get("pr_number", "")) cedar_policies = inp.get("cedar_policies") or [] + # Registry assets (#246) resolved by the orchestrator; forwarded verbatim to + # the pipeline, which applies the per-kind loaders (mcp_server → .mcp.json). + resolved_assets = inp.get("resolved_assets") or [] # Cedar HITL (§7.3) — per-task approval defaults + seeded allowlist. # Both are forwarded verbatim to the pipeline; the engine # validates shape at construction time and raises on bad input. @@ -642,6 +647,7 @@ def _extract_invocation_params(inp: dict, request: Request) -> dict: "branch_name": branch_name, "pr_number": pr_number, "cedar_policies": cedar_policies, + "resolved_assets": resolved_assets, "approval_timeout_s": approval_timeout_s, "initial_approvals": initial_approvals, "initial_approval_gate_count": initial_approval_gate_count, diff --git a/agent/tests/test_registry_loader.py b/agent/tests/test_registry_loader.py new file mode 100644 index 000000000..ae351efee --- /dev/null +++ b/agent/tests/test_registry_loader.py @@ -0,0 +1,151 @@ +"""Unit tests for registry.loader — merging resolved mcp_server assets (#246).""" + +from __future__ import annotations + +import json + +from registry.loader import ( + apply_mcp_assets, + apply_resolved_assets, + build_skill_prompt_fragment, +) + + +def _read_mcp(repo_dir) -> dict: + with open(repo_dir / ".mcp.json", encoding="utf-8") as f: + return json.load(f) + + +def _mcp_asset(namespace: str, name: str, version: str, runtime: dict) -> dict: + return { + "kind": "mcp_server", + "namespace": namespace, + "name": name, + "version": version, + "runtime": runtime, + } + + +class TestApplyMcpAssets: + def test_writes_new_mcp_json(self, tmp_path): + runtime = {"transport": "http", "url": "https://mcp.example.com/sse"} + n = apply_mcp_assets(str(tmp_path), [_mcp_asset("acme", "pdf-tools", "1.0.0", runtime)]) + assert n == 1 + merged = _read_mcp(tmp_path) + assert merged["mcpServers"]["acme__pdf_tools"] == runtime + + def test_preserves_existing_servers(self, tmp_path): + existing = {"mcpServers": {"other": {"command": "/usr/bin/x"}}} + (tmp_path / ".mcp.json").write_text(json.dumps(existing)) + n = apply_mcp_assets( + str(tmp_path), + [_mcp_asset("acme", "weather", "2.1.0", {"transport": "sse", "url": "https://w"})], + ) + assert n == 1 + merged = _read_mcp(tmp_path) + assert merged["mcpServers"]["other"]["command"] == "/usr/bin/x" + assert "acme__weather" in merged["mcpServers"] + + def test_merges_multiple_servers(self, tmp_path): + assets = [ + _mcp_asset("acme", "a", "1.0.0", {"transport": "http", "url": "https://a"}), + _mcp_asset("acme", "b", "1.0.0", {"transport": "http", "url": "https://b"}), + ] + n = apply_mcp_assets(str(tmp_path), assets) + assert n == 2 + merged = _read_mcp(tmp_path) + assert set(merged["mcpServers"]) == {"acme__a", "acme__b"} + + def test_ignores_non_mcp_kinds(self, tmp_path): + assets = [ + { + "kind": "cedar_policy_module", + "namespace": "acme", + "name": "p", + "version": "1.0.0", + "runtime": {"cedar_text": "permit(...);"}, + }, + ] + n = apply_mcp_assets(str(tmp_path), assets) + assert n == 0 + assert not (tmp_path / ".mcp.json").exists() + + def test_skips_empty_runtime(self, tmp_path): + n = apply_mcp_assets(str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", {})]) + assert n == 0 + assert not (tmp_path / ".mcp.json").exists() + + def test_no_op_on_missing_repo_dir(self): + asset = _mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "u"}) + n = apply_mcp_assets("/nonexistent/dir", [asset]) + assert n == 0 + + def test_malformed_existing_treated_as_absent(self, tmp_path): + (tmp_path / ".mcp.json").write_text("{ not valid json") + runtime = {"transport": "http", "url": "https://x"} + n = apply_mcp_assets(str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", runtime)]) + assert n == 1 + assert _read_mcp(tmp_path)["mcpServers"]["acme__x"] == runtime + + +def _skill_asset(namespace: str, name: str, runtime: dict) -> dict: + return { + "kind": "skill", + "namespace": namespace, + "name": name, + "version": "1.0.0", + "runtime": runtime, + } + + +class TestBuildSkillPromptFragment: + def test_empty_when_no_skills(self): + assert build_skill_prompt_fragment([]) == "" + mcp = _mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "u"}) + assert build_skill_prompt_fragment([mcp]) == "" + + def test_appends_fragment_with_heading(self): + out = build_skill_prompt_fragment( + [_skill_asset("acme", "research", {"prompt_fragment": "Summarize findings."})] + ) + assert "## Skills" in out + assert "### Skill: acme/research" in out + assert "Summarize findings." in out + + def test_includes_tool_hints(self): + runtime = {"prompt_fragment": "Do X.", "tool_hints": ["Bash", "Edit"]} + out = build_skill_prompt_fragment([_skill_asset("acme", "r", runtime)]) + assert "Bash, Edit" in out + + def test_concatenates_multiple_in_order(self): + out = build_skill_prompt_fragment( + [ + _skill_asset("acme", "a", {"prompt_fragment": "First."}), + _skill_asset("acme", "b", {"prompt_fragment": "Second."}), + ] + ) + assert out.index("First.") < out.index("Second.") + + def test_skips_blank_or_invalid_runtime(self): + blank = _skill_asset("acme", "a", {"prompt_fragment": " "}) + assert build_skill_prompt_fragment([blank]) == "" + assert build_skill_prompt_fragment([_skill_asset("acme", "a", {})]) == "" + + +class TestApplyResolvedAssets: + def test_empty_is_noop(self, tmp_path): + apply_resolved_assets(str(tmp_path), []) + assert not (tmp_path / ".mcp.json").exists() + + def test_dispatches_mcp(self, tmp_path): + apply_resolved_assets( + str(tmp_path), + [_mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "https://x"})], + ) + assert (tmp_path / ".mcp.json").exists() + + def test_skill_and_cedar_do_not_touch_mcp_json(self, tmp_path): + # apply_resolved_assets only handles on-disk kinds (mcp_server). Skills + # and cedar modules are applied elsewhere, so no .mcp.json is written. + apply_resolved_assets(str(tmp_path), [_skill_asset("acme", "r", {"prompt_fragment": "X."})]) + assert not (tmp_path / ".mcp.json").exists() diff --git a/cdk/src/constructs/blueprint.ts b/cdk/src/constructs/blueprint.ts index 5ac64ac1d..73aff979a 100644 --- a/cdk/src/constructs/blueprint.ts +++ b/cdk/src/constructs/blueprint.ts @@ -26,6 +26,7 @@ import { Construct, IValidation } from 'constructs'; // the JSON directly rather than re-using ``handlers/shared/types.ts`` so // the construct layer stays decoupled from runtime-side types. import sharedConstants from '../../../contracts/constants.json'; +import { parseRef } from '../handlers/shared/registry/ref'; const REPO_PATTERN = /^[a-zA-Z0-9._-]+\/[a-zA-Z0-9._-]+$/; const DOMAIN_PATTERN = /^(\*\.)?[a-z0-9]([a-z0-9-]*[a-z0-9])?(\.[a-z0-9]([a-z0-9-]*[a-z0-9])?)*$/; @@ -149,6 +150,21 @@ export interface BlueprintProps { */ readonly egressAllowlist?: string[]; }; + + /** + * Registry assets (#246) this repo pins. Each entry is a strict + * ``registry://kind/namespace/name@constraint`` ref, validated at synth. The + * orchestrator resolves the refs at task start and threads the resolved bundle + * into the agent payload; an unresolvable ref fails the task (fail-closed). + */ + readonly assets?: { + /** MCP servers merged into the agent's ``.mcp.json`` (PR 2). */ + readonly mcpServers?: string[]; + /** Cedar policy modules concatenated into the agent's cedar_policies (PR 3). */ + readonly cedarPolicyModules?: string[]; + /** Skills whose prompt fragments are appended to the system prompt (PR 3). */ + readonly skills?: string[]; + }; } /** @@ -183,12 +199,26 @@ export class Blueprint extends Construct { */ public readonly approvalGateCap?: number; + /** + * Registry ``registry://`` refs for MCP servers (#246), exposed for inspection. + */ + public readonly mcpServerRefs: readonly string[]; + + /** Registry ``registry://`` refs for Cedar policy modules (#246). */ + public readonly cedarPolicyModuleRefs: readonly string[]; + + /** Registry ``registry://`` refs for skills (#246). */ + public readonly skillRefs: readonly string[]; + constructor(scope: Construct, id: string, props: BlueprintProps) { super(scope, id); this.egressAllowlist = [...(props.networking?.egressAllowlist ?? [])]; this.cedarPolicies = [...(props.security?.cedarPolicies ?? [])]; this.approvalGateCap = props.security?.approvalGateCap; + this.mcpServerRefs = [...(props.assets?.mcpServers ?? [])]; + this.cedarPolicyModuleRefs = [...(props.assets?.cedarPolicyModules ?? [])]; + this.skillRefs = [...(props.assets?.skills ?? [])]; // Chunk 7c: emit a synth-time info annotation when the blueprint did // not configure an override so operators see a signal that this repo @@ -207,6 +237,9 @@ export class Blueprint extends Construct { this.node.addValidation(new RepoFormatValidation(props.repo)); this.node.addValidation(new DomainFormatValidation(this.egressAllowlist)); this.node.addValidation(new ApprovalGateCapValidation(this.approvalGateCap)); + this.node.addValidation(new RegistryRefValidation('assets.mcpServers', this.mcpServerRefs)); + this.node.addValidation(new RegistryRefValidation('assets.cedarPolicyModules', this.cedarPolicyModuleRefs)); + this.node.addValidation(new RegistryRefValidation('assets.skills', this.skillRefs)); const now = new Date().toISOString(); @@ -248,6 +281,15 @@ export class Blueprint extends Construct { if (this.approvalGateCap !== undefined) { item.approval_gate_cap = { N: String(this.approvalGateCap) }; } + if (this.mcpServerRefs.length > 0) { + item.mcp_servers = { L: this.mcpServerRefs.map(r => ({ S: r })) }; + } + if (this.cedarPolicyModuleRefs.length > 0) { + item.cedar_policy_modules = { L: this.cedarPolicyModuleRefs.map(r => ({ S: r })) }; + } + if (this.skillRefs.length > 0) { + item.skills = { L: this.skillRefs.map(r => ({ S: r })) }; + } new cr.AwsCustomResource(this, 'RepoConfigCR', { timeout: Duration.minutes(REPO_CONFIG_CR_TIMEOUT_MINUTES), @@ -320,6 +362,11 @@ export class Blueprint extends Construct { if (this.egressAllowlist.length > 0) fields.push(', #egress_allowlist = :egress_allowlist'); if (this.cedarPolicies.length > 0) fields.push(', #cedar_policies = :cedar_policies'); if (this.approvalGateCap !== undefined) fields.push(', #approval_gate_cap = :approval_gate_cap'); + // Registry asset refs (#246) — must mirror onCreate's item, else a redeploy + // of an already-onboarded repo silently drops asset-ref changes. + if (this.mcpServerRefs.length > 0) fields.push(', #mcp_servers = :mcp_servers'); + if (this.cedarPolicyModuleRefs.length > 0) fields.push(', #cedar_policy_modules = :cedar_policy_modules'); + if (this.skillRefs.length > 0) fields.push(', #skills = :skills'); return fields.join(''); } @@ -335,6 +382,9 @@ export class Blueprint extends Construct { if (this.egressAllowlist.length > 0) names['#egress_allowlist'] = 'egress_allowlist'; if (this.cedarPolicies.length > 0) names['#cedar_policies'] = 'cedar_policies'; if (this.approvalGateCap !== undefined) names['#approval_gate_cap'] = 'approval_gate_cap'; + if (this.mcpServerRefs.length > 0) names['#mcp_servers'] = 'mcp_servers'; + if (this.cedarPolicyModuleRefs.length > 0) names['#cedar_policy_modules'] = 'cedar_policy_modules'; + if (this.skillRefs.length > 0) names['#skills'] = 'skills'; return names; } @@ -350,6 +400,9 @@ export class Blueprint extends Construct { if (this.egressAllowlist.length > 0) values[':egress_allowlist'] = { L: this.egressAllowlist.map(d => ({ S: d })) }; if (this.cedarPolicies.length > 0) values[':cedar_policies'] = { L: this.cedarPolicies.map(p => ({ S: p })) }; if (this.approvalGateCap !== undefined) values[':approval_gate_cap'] = { N: String(this.approvalGateCap) }; + if (this.mcpServerRefs.length > 0) values[':mcp_servers'] = { L: this.mcpServerRefs.map(r => ({ S: r })) }; + if (this.cedarPolicyModuleRefs.length > 0) values[':cedar_policy_modules'] = { L: this.cedarPolicyModuleRefs.map(r => ({ S: r })) }; + if (this.skillRefs.length > 0) values[':skills'] = { L: this.skillRefs.map(r => ({ S: r })) }; return values; } } @@ -411,3 +464,23 @@ class ApprovalGateCapValidation implements IValidation { return []; } } + +/** + * Registry (#246) — validates each ``registry://`` asset ref against the strict + * grammar at synth, so a floating or malformed pin cannot deploy and then fail + * every task at resolve time. Uses the same ``parseRef`` the resolver enforces. + */ +class RegistryRefValidation implements IValidation { + constructor(private readonly field: string, private readonly refs: readonly string[]) {} + + public validate(): string[] { + const errors: string[] = []; + for (const ref of this.refs) { + const result = parseRef(ref); + if (!result.ok) { + errors.push(`Invalid ${this.field} ref '${ref}': ${result.reason} — ${result.message}`); + } + } + return errors; + } +} diff --git a/cdk/src/constructs/task-orchestrator.ts b/cdk/src/constructs/task-orchestrator.ts index a638b369f..f9370dfbc 100644 --- a/cdk/src/constructs/task-orchestrator.ts +++ b/cdk/src/constructs/task-orchestrator.ts @@ -18,7 +18,7 @@ */ import * as path from 'path'; -import { Duration, Stack } from 'aws-cdk-lib'; +import { ArnFormat, Duration, Stack } from 'aws-cdk-lib'; import * as cloudwatch from 'aws-cdk-lib/aws-cloudwatch'; import * as dynamodb from 'aws-cdk-lib/aws-dynamodb'; import * as iam from 'aws-cdk-lib/aws-iam'; @@ -174,6 +174,13 @@ export interface TaskOrchestratorProps { * ReadWrite grants for URL fetch/screen/upload during hydration. */ readonly attachmentsBucket?: s3.IBucket; + + /** + * AgentCore registry id (#246). When provided, the orchestrator resolves the + * Blueprint's ``registry://`` asset refs at task start and threads the bundle + * into the agent payload. Requires bedrock-agentcore registry read actions. + */ + readonly agentRegistryId?: string; } /** @@ -277,6 +284,7 @@ export class TaskOrchestrator extends Construct { // from at finalize); the ECS strategy reads this to build the S3 URI. ...(props.ecsPayloadBucket && { ECS_PAYLOAD_BUCKET: props.ecsPayloadBucket.bucketName }), ...(props.attachmentsBucket && { ATTACHMENTS_BUCKET_NAME: props.attachmentsBucket.bucketName }), + ...(props.agentRegistryId && { AGENT_REGISTRY_ID: props.agentRegistryId }), }, bundling: orchestratorBundling, }); @@ -337,6 +345,32 @@ export class TaskOrchestrator extends Construct { resources: runtimeResources, })); + // Registry (#246): read-only access so the orchestrator can resolve the + // Blueprint's registry:// asset refs at task start. Record ids are + // server-assigned, so the record ARN is a wildcard under the registry. + if (props.agentRegistryId) { + this.fn.addToRolePolicy(new iam.PolicyStatement({ + actions: [ + 'bedrock-agentcore:GetRegistryRecord', + 'bedrock-agentcore:ListRegistryRecords', + ], + resources: [ + Stack.of(this).formatArn({ + service: 'bedrock-agentcore', + resource: 'registry', + resourceName: '*', + arnFormat: ArnFormat.SLASH_RESOURCE_NAME, + }), + Stack.of(this).formatArn({ + service: 'bedrock-agentcore', + resource: 'registry', + resourceName: '*/record/*', + arnFormat: ArnFormat.SLASH_RESOURCE_NAME, + }), + ], + })); + } + // ECS compute strategy permissions (only when ECS is configured) if (props.ecsConfig) { this.fn.addToRolePolicy(new iam.PolicyStatement({ @@ -415,7 +449,7 @@ export class TaskOrchestrator extends Construct { }, { id: 'AwsSolutions-IAM5', - reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData; AgentCore runtime/* required for sub-resource invocation; Secrets Manager wildcards generated by CDK grantRead; AgentCore Memory wildcards generated by CDK grantRead/grantWrite; ECS RunTask/DescribeTasks/StopTask conditioned on cluster ARN; iam:PassRole scoped to ECS task/execution roles and conditioned on ecs-tasks.amazonaws.com', + reason: 'DynamoDB index/* wildcards generated by CDK grantReadWriteData; AgentCore runtime/* required for sub-resource invocation; Secrets Manager wildcards generated by CDK grantRead; AgentCore Memory wildcards generated by CDK grantRead/grantWrite; ECS RunTask/DescribeTasks/StopTask conditioned on cluster ARN; iam:PassRole scoped to ECS task/execution roles and conditioned on ecs-tasks.amazonaws.com; AgentCore registry/* + registry/*/record/* wildcards because record ids are server-assigned (#246)', }, ], true); } diff --git a/cdk/src/handlers/shared/orchestrator.ts b/cdk/src/handlers/shared/orchestrator.ts index 7308c4532..57063c156 100644 --- a/cdk/src/handlers/shared/orchestrator.ts +++ b/cdk/src/handlers/shared/orchestrator.ts @@ -26,6 +26,9 @@ import { logger, type Logger } from './logger'; import { writeMinimalEpisode } from './memory'; import { coerceNumericOrNull } from './numeric'; import { computePromptVersion } from './prompt-version'; +import { makeRegistryClient } from './registry/factory'; +import { parseRef } from './registry/ref'; +import { RegistryResolutionError, type ResolvedAsset } from './registry/types'; import { loadRepoConfig, type BlueprintConfig, type ComputeType } from './repo-config'; import { resolveUrlAttachments } from './resolve-url-attachments'; import { APPROVAL_GATE_CAP_MAX, APPROVAL_GATE_CAP_MIN, type AgentAttachmentPayload, type AttachmentRecord, type TaskRecord } from './types'; @@ -279,9 +282,49 @@ export async function loadBlueprintConfig(task: TaskRecord): Promise { + const refs = [ + ...(blueprintConfig?.mcp_servers ?? []), + ...(blueprintConfig?.cedar_policy_modules ?? []), + ...(blueprintConfig?.skills ?? []), + ]; + if (refs.length === 0) return []; + + const client = makeRegistryClient(); + const resolved: ResolvedAsset[] = []; + for (const ref of refs) { + const parsed = parseRef(ref); + if (!parsed.ok) { + throw new RegistryResolutionError(parsed.reason, ref, parsed.message); + } + const asset = await client.resolve(parsed.ref); + if (asset.warnings.length > 0) { + log.warn('Registry asset resolved with warnings', { ref, warnings: asset.warnings }); + } + resolved.push(asset); + } + log.info('Resolved registry assets', { count: resolved.length }); + return resolved; +} + /** * Map passed AttachmentRecords into the payload shape the agent runtime expects. * Only includes attachments that passed screening (others are already rejected). @@ -523,6 +566,35 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B ? resolveAttachmentPayloads(resolvedAttachments, Number(process.env.USER_PROMPT_TOKEN_BUDGET ?? '100000')) : []; + // Resolve registry assets (#246). Fail-closed: an unresolved ref throws here + // and the orchestrator transitions the task to FAILED. The audit triple is + // stamped on the TaskRecord; the runtime bundle rides in the payload. + const resolvedAssets = await resolveRegistryAssets(blueprintConfig, log); + if (resolvedAssets.length > 0) { + await ddb.send(new UpdateCommand({ + TableName: TABLE_NAME, + Key: { task_id: task.task_id }, + UpdateExpression: 'SET #ra = :ra, #ua = :now', + ExpressionAttributeNames: { '#ra': 'resolved_assets', '#ua': 'updated_at' }, + ExpressionAttributeValues: { + ':ra': resolvedAssets.map((a) => ({ kind: a.kind, id: `${a.namespace}/${a.name}`, version: a.version })), + ':now': new Date().toISOString(), + }, + })); + } + + // Registry cedar_policy_module assets (#246, PR 3) reach the agent through the + // SAME cedar_policies payload field as inline blueprint policies, so they are + // byte-identical from the PolicyEngine's view (the cedar-parity contract holds + // by construction). Inline blueprint policies come first, then resolved modules. + const cedarText = [ + ...(blueprintConfig?.cedar_policies ?? []), + ...resolvedAssets + .filter((a) => a.kind === 'cedar_policy_module') + .map((a) => (a.runtime as { cedar_text?: string }).cedar_text) + .filter((t): t is string => typeof t === 'string' && t.length > 0), + ]; + const payload: Record = { repo_url: task.repo, task_id: task.task_id, @@ -548,7 +620,14 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B ...(task.trace === true && { trace: true }), ...(blueprintConfig?.model_id && { model_id: blueprintConfig.model_id }), ...(blueprintConfig?.system_prompt_overrides && { system_prompt_overrides: blueprintConfig.system_prompt_overrides }), - ...(blueprintConfig?.cedar_policies && blueprintConfig.cedar_policies.length > 0 && { cedar_policies: blueprintConfig.cedar_policies }), + ...(cedarText.length > 0 && { cedar_policies: cedarText }), + // Registry (#246): the resolved runtime bundle the agent's loaders apply + // (MCP servers merged into .mcp.json in PR 2; cedar/skills in PR 3). + ...(resolvedAssets.length > 0 && { + resolved_assets: resolvedAssets.map((a) => ({ + kind: a.kind, namespace: a.namespace, name: a.name, version: a.version, runtime: a.runtime, + })), + }), // Cedar HITL: the agent's PreToolUse hook uses this to compute // the maxLifetime ceiling on per-gate approval timeouts (§6.5). // Stamped at HYDRATING → RUNNING transition time so the clock diff --git a/cdk/src/handlers/shared/repo-config.ts b/cdk/src/handlers/shared/repo-config.ts index 1753d5685..ea772a1ca 100644 --- a/cdk/src/handlers/shared/repo-config.ts +++ b/cdk/src/handlers/shared/repo-config.ts @@ -52,6 +52,18 @@ export interface RepoConfig { * path falls back to the platform default of 50. */ readonly approval_gate_cap?: number; + /** + * Registry (#246) ``registry://`` refs for MCP servers pinned by the + * blueprint. Resolved by the orchestrator at task start and merged into the + * agent's ``.mcp.json``. + */ + readonly mcp_servers?: string[]; + /** Registry (#246) Cedar policy module refs; resolved cedar_text is merged + * into the ``cedar_policies`` payload. */ + readonly cedar_policy_modules?: string[]; + /** Registry (#246) skill refs; resolved prompt fragments append to the + * system prompt. */ + readonly skills?: string[]; } /** @@ -77,6 +89,15 @@ export interface BlueprintConfig { * field is informational for the runtime path. */ readonly approval_gate_cap?: number; + /** + * Registry (#246) MCP server ``registry://`` refs surfaced from RepoConfig so + * the orchestrator can resolve + merge them into the agent payload. + */ + readonly mcp_servers?: string[]; + /** Registry (#246) Cedar policy module refs surfaced from RepoConfig. */ + readonly cedar_policy_modules?: string[]; + /** Registry (#246) skill refs surfaced from RepoConfig. */ + readonly skills?: string[]; } const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({})); diff --git a/cdk/src/stacks/agent.ts b/cdk/src/stacks/agent.ts index 045d6ae36..7b29bedf9 100644 --- a/cdk/src/stacks/agent.ts +++ b/cdk/src/stacks/agent.ts @@ -176,6 +176,23 @@ export class AgentStack extends Stack { const blueprints = [agentPluginsBlueprint]; + // Optional per-repo blueprint pinning registry assets (#246), opt-in via + // context/env so it does not hardcode a specific fork for other contributors. + // Set ``forkBlueprintRepo`` (e.g. ``--context forkBlueprintRepo=owner/repo``) + // to onboard a repo with the AWS Knowledge MCP asset pinned. + const forkBlueprintRepo = process.env.FORK_BLUEPRINT_REPO ?? this.node.tryGetContext('forkBlueprintRepo'); + if (forkBlueprintRepo) { + blueprints.push(new Blueprint(this, 'ForkBlueprint', { + repo: forkBlueprintRepo, + repoTable: repoTable.table, + assets: { + mcpServers: ['registry://mcp_server/acme/aws-knowledge@^1.0.0'], + cedarPolicyModules: ['registry://cedar_policy_module/acme/guard@^1.0.0'], + skills: ['registry://skill/acme/readme-helper@^1.0.0'], + }, + })); + } + // The AwsCustomResource singleton Lambda used by Blueprint constructs NagSuppressions.addResourceSuppressionsByPath(this, [ `${this.stackName}/AWS679f53fac002430cb0da5b7982bd2287/ServiceRole/Resource`, @@ -693,6 +710,7 @@ export class AgentStack extends Stack { guardrailId: inputGuardrail.guardrailId, guardrailVersion: inputGuardrail.guardrailVersion, attachmentsBucket: attachmentsBucket.bucket, + agentRegistryId: agentRegistry.registryId, // K12: route ``compute_type: 'ecs'`` repos to the Fargate cluster above — // only when the cluster was synthesized (deploy --context compute_type=ecs). ...(ecsCluster && { diff --git a/cdk/test/constructs/blueprint.test.ts b/cdk/test/constructs/blueprint.test.ts index 57eca5096..b1ac03eab 100644 --- a/cdk/test/constructs/blueprint.test.ts +++ b/cdk/test/constructs/blueprint.test.ts @@ -300,6 +300,73 @@ describe('Blueprint construct', () => { expect(serialized).toContain('#cedar_policies'); }); + // --- Registry asset refs (#246) --- + + test('maps registry asset refs to DynamoDB lists', () => { + const { template } = createStack({ + assets: { + mcpServers: ['registry://mcp_server/acme/pdf-tools@^1.4.1'], + cedarPolicyModules: ['registry://cedar_policy_module/acme/force-push@~2.0.0'], + skills: ['registry://skill/acme/research@1.0.0'], + }, + }); + const serialized = getCreateJoinParts(template).join(''); + expect(serialized).toContain('"mcp_servers":{"L":[{"S":"registry://mcp_server/acme/pdf-tools@^1.4.1"}]}'); + expect(serialized).toContain('"cedar_policy_modules":{"L":[{"S":"registry://cedar_policy_module/acme/force-push@~2.0.0"}]}'); + expect(serialized).toContain('"skills":{"L":[{"S":"registry://skill/acme/research@1.0.0"}]}'); + }); + + test('omits asset columns when no assets are pinned', () => { + const serialized = getCreateJoinParts(createStack().template).join(''); + expect(serialized).not.toContain('mcp_servers'); + expect(serialized).not.toContain('cedar_policy_modules'); + expect(serialized).not.toContain('"skills"'); + }); + + test('onUpdate also writes asset refs (redeploy of an onboarded repo must not drop them)', () => { + const { template } = createStack({ + assets: { + mcpServers: ['registry://mcp_server/acme/pdf-tools@^1.4.1'], + cedarPolicyModules: ['registry://cedar_policy_module/acme/force-push@~2.0.0'], + skills: ['registry://skill/acme/research@1.0.0'], + }, + }); + const serialized = getUpdateJoinParts(template).join(''); + // Regression guard (#246): the onUpdate UpdateExpression previously omitted + // the three asset-ref columns, so a redeploy silently dropped them. + expect(serialized).toContain('#mcp_servers'); + expect(serialized).toContain('#cedar_policy_modules'); + expect(serialized).toContain('#skills'); + }); + + test('rejects a floating asset ref at synth', () => { + const app = new App(); + const stack = new Stack(app, 'TestStack'); + const repoTable = new dynamodb.Table(stack, 'RepoTable', { + partitionKey: { name: 'repo', type: dynamodb.AttributeType.STRING }, + }); + new Blueprint(stack, 'Blueprint', { + repo: 'org/my-repo', + repoTable, + assets: { mcpServers: ['registry://mcp_server/acme/pdf-tools'] }, + }); + expect(() => Template.fromStack(stack)).toThrow(/Invalid assets.mcpServers ref.*INVALID_REGISTRY_REF/); + }); + + test('rejects a malformed constraint on a skill ref at synth', () => { + const app = new App(); + const stack = new Stack(app, 'TestStack'); + const repoTable = new dynamodb.Table(stack, 'RepoTable', { + partitionKey: { name: 'repo', type: dynamodb.AttributeType.STRING }, + }); + new Blueprint(stack, 'Blueprint', { + repo: 'org/my-repo', + repoTable, + assets: { skills: ['registry://skill/acme/research@latest'] }, + }); + expect(() => Template.fromStack(stack)).toThrow(/Invalid assets.skills ref.*INVALID_CONSTRAINT/); + }); + // --- Chunk 7b: security.approvalGateCap --------------------------------- test('exposes approvalGateCap as public property when configured', () => { diff --git a/cdk/test/handlers/shared/registry-orchestrator.test.ts b/cdk/test/handlers/shared/registry-orchestrator.test.ts new file mode 100644 index 000000000..fc59a8851 --- /dev/null +++ b/cdk/test/handlers/shared/registry-orchestrator.test.ts @@ -0,0 +1,141 @@ +/** + * MIT No Attribution + * + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * the Software without restriction, including without limitation the rights to + * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of + * the Software, and to permit persons to whom the Software is furnished to do so. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + */ + +/** + * E2E-ish coverage of the orchestrator registry resolve-step (#246, PR 2): + * given a Blueprint's ``mcp_servers`` refs, ``resolveRegistryAssets`` resolves + * each via the RegistryClient and is fail-closed on a bad ref / resolution + * failure. This is the seam the full task path calls before assembling the + * agent payload; the payload/stamping wiring around it is exercised here by + * asserting the returned bundle shape the orchestrator threads through. + */ + +import { resolveRegistryAssets } from '../../../src/handlers/shared/orchestrator'; +import { RegistryResolutionError, type ResolvedAsset } from '../../../src/handlers/shared/registry/types'; +import type { BlueprintConfig } from '../../../src/handlers/shared/repo-config'; + +// Standalone mock fn (not a method on an object) so `.not.toHaveBeenCalled()` +// doesn't trip @typescript-eslint/unbound-method — matches the repo pattern. +const mockResolve = jest.fn(); +jest.mock('../../../src/handlers/shared/registry/factory', () => { + const actual = jest.requireActual('../../../src/handlers/shared/registry/factory'); + return { ...actual, makeRegistryClient: () => ({ resolve: mockResolve }) }; +}); + +const log = { info: jest.fn(), warn: jest.fn(), error: jest.fn() } as never; + +const asset = (over: Partial = {}): ResolvedAsset => ({ + kind: 'mcp_server', + namespace: 'acme', + name: 'pdf-tools', + version: '1.4.1', + runtime: { transport: 'http', url: 'https://mcp.example.com/sse' } as never, + warnings: [], + ...over, +}); + +const bp = (refs: Partial> = {}): BlueprintConfig => ({ + compute_type: 'agentcore', + runtime_arn: 'arn:aws:bedrock-agentcore:us-east-1:1:runtime/r', + ...refs, +}); + +beforeEach(() => jest.clearAllMocks()); + +describe('resolveRegistryAssets', () => { + test('returns [] when the blueprint pins no assets', async () => { + expect(await resolveRegistryAssets(bp(), log)).toEqual([]); + expect(await resolveRegistryAssets(undefined, log)).toEqual([]); + expect(mockResolve).not.toHaveBeenCalled(); + }); + + test('resolves each ref into the bundle the orchestrator threads', async () => { + mockResolve.mockResolvedValue(asset()); + const result = await resolveRegistryAssets( + bp({ mcp_servers: ['registry://mcp_server/acme/pdf-tools@^1.4.1'] }), + log, + ); + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ kind: 'mcp_server', namespace: 'acme', name: 'pdf-tools', version: '1.4.1' }); + expect(result[0].runtime).toMatchObject({ transport: 'http' }); + }); + + test('resolves multiple refs in order', async () => { + mockResolve + .mockResolvedValueOnce(asset({ name: 'a', version: '1.0.0' })) + .mockResolvedValueOnce(asset({ name: 'b', version: '2.0.0' })); + const result = await resolveRegistryAssets( + bp({ mcp_servers: ['registry://mcp_server/acme/a@^1.0.0', 'registry://mcp_server/acme/b@^2.0.0'] }), + log, + ); + expect(result.map((a) => a.name)).toEqual(['a', 'b']); + }); + + test('fail-closed on a malformed ref (never calls resolve)', async () => { + await expect( + resolveRegistryAssets(bp({ mcp_servers: ['registry://mcp_server/acme/pdf-tools'] }), log), + ).rejects.toBeInstanceOf(RegistryResolutionError); + expect(mockResolve).not.toHaveBeenCalled(); + }); + + test('fail-closed when the client cannot resolve a version', async () => { + mockResolve.mockRejectedValue( + new RegistryResolutionError('NO_MATCHING_VERSION', 'r', 'none'), + ); + await expect( + resolveRegistryAssets(bp({ mcp_servers: ['registry://mcp_server/acme/pdf-tools@^9.9.9'] }), log), + ).rejects.toMatchObject({ reason: 'NO_MATCHING_VERSION' }); + }); + + test('a DEPRECATED asset resolves but is logged as a warning', async () => { + mockResolve.mockResolvedValue(asset({ warnings: ['DEPRECATED'] })); + const result = await resolveRegistryAssets( + bp({ mcp_servers: ['registry://mcp_server/acme/pdf-tools@1.4.1'] }), + log, + ); + expect(result).toHaveLength(1); + expect(log.warn).toHaveBeenCalled(); + }); + + test('resolves cedar_policy_module + skill refs alongside mcp (PR 3)', async () => { + mockResolve + .mockResolvedValueOnce(asset({ kind: 'mcp_server', name: 'pdf-tools' })) + .mockResolvedValueOnce(asset({ + kind: 'cedar_policy_module', + name: 'force-push', + runtime: { cedar_text: 'forbid(principal, action, resource);' } as never, + })) + .mockResolvedValueOnce(asset({ + kind: 'skill', + name: 'research', + runtime: { prompt_fragment: 'Summarize.' } as never, + })); + const result = await resolveRegistryAssets( + bp({ + mcp_servers: ['registry://mcp_server/acme/pdf-tools@^1.4.1'], + cedar_policy_modules: ['registry://cedar_policy_module/acme/force-push@^1.0.0'], + skills: ['registry://skill/acme/research@^1.0.0'], + }), + log, + ); + expect(result.map((a) => a.kind)).toEqual(['mcp_server', 'cedar_policy_module', 'skill']); + expect((result[1].runtime as { cedar_text: string }).cedar_text).toContain('forbid'); + expect((result[2].runtime as { prompt_fragment: string }).prompt_fragment).toBe('Summarize.'); + }); +});