From 4711ab9a28191afa99a6baeb5544bc89a3a5aed5 Mon Sep 17 00:00:00 2001 From: bgagent Date: Mon, 27 Jul 2026 13:17:42 -0400 Subject: [PATCH 1/3] 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 | 85 +++++++++- 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, 776 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 0a3c4d0c3..556b4e58e 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 @@ -228,6 +228,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 human-in-the-loop approvals. 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 6742018e2..5cb5687f9 100644 --- a/agent/src/pipeline.py +++ b/agent/src/pipeline.py @@ -833,6 +833,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. @@ -887,6 +888,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). @@ -1147,6 +1153,13 @@ def _on_trace_truncated(max_bytes: int, first_dropped: int) -> None: # matches Jira's own entry. strip_linear_mcp_servers(setup.repo_dir) + # 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 3582575d6..d95992a81 100644 --- a/agent/src/prompt_builder.py +++ b/agent/src/prompt_builder.py @@ -80,6 +80,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 @@ -113,6 +120,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 831696555..a6be69f42 100644 --- a/agent/src/server.py +++ b/agent/src/server.py @@ -412,6 +412,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 @@ -501,6 +502,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: @@ -555,6 +557,9 @@ def _extract_invocation_params(inp: dict, request: Request) -> dict: merge_branches_raw = inp.get("merge_branches") or [] merge_branches = [b for b in merge_branches_raw if isinstance(b, str)] 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. @@ -665,6 +670,7 @@ def _extract_invocation_params(inp: dict, request: Request) -> dict: "base_branch": base_branch, "merge_branches": merge_branches, "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 354634424..751d52c3c 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])?)*$/; @@ -170,6 +171,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[]; + }; } /** @@ -204,12 +220,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 @@ -228,6 +258,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(); @@ -275,6 +308,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), @@ -349,6 +391,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(''); } @@ -366,6 +413,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; } @@ -383,6 +433,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; } } @@ -444,3 +497,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 fcaf4d282..e21307e4e 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'; @@ -265,6 +265,13 @@ export interface TaskOrchestratorProps { */ readonly payloadBucket: 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; } /** @@ -395,6 +402,7 @@ export class TaskOrchestrator extends Construct { }), }), ...(props.attachmentsBucket && { ATTACHMENTS_BUCKET_NAME: props.attachmentsBucket.bucketName }), + ...(props.agentRegistryId && { AGENT_REGISTRY_ID: props.agentRegistryId }), }, bundling: orchestratorBundling, }); @@ -465,6 +473,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({ @@ -628,7 +662,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; S3 object/* wildcard from CDK grantPut on the dedicated MicroVM payload bucket; MicroVM lifecycle actions (RunMicrovm/GetMicrovm/TerminateMicrovm) are scoped to the single platform MicroVM image ARN plus a :* version-suffix sibling (every one of them authorizes against the image resource, not the per-session instance; no account-wide wildcard is used); lambda:PassNetworkConnector requires Resource:* because the action supports no resource-level permissions and the AWS-managed connectors live outside this account; iam:PassRole is scoped to the MicroVM execution role and conditioned on lambda.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; S3 object/* wildcard from CDK grantPut on the dedicated MicroVM payload bucket; MicroVM lifecycle actions (RunMicrovm/GetMicrovm/TerminateMicrovm) are scoped to the single platform MicroVM image ARN plus a :* version-suffix sibling (every one of them authorizes against the image resource, not the per-session instance; no account-wide wildcard is used); lambda:PassNetworkConnector requires Resource:* because the action supports no resource-level permissions and the AWS-managed connectors live outside this account; iam:PassRole is scoped to the MicroVM execution role and conditioned on lambda.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 f91950f1b..77bce1786 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'; @@ -502,9 +505,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). @@ -746,6 +789,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, @@ -785,7 +857,18 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B // build-regression gating actually runs the repo's real command. ...(blueprintConfig?.build_command && { build_command: blueprintConfig.build_command }), ...(blueprintConfig?.lint_command && { lint_command: blueprintConfig.lint_command }), - ...(blueprintConfig?.cedar_policies && blueprintConfig.cedar_policies.length > 0 && { cedar_policies: blueprintConfig.cedar_policies }), + // cedarText is inline blueprint policies ++ resolved registry + // cedar_policy_module text (#246), so it supersedes the raw + // blueprintConfig.cedar_policies — byte-identical to inline when no + // registry cedar is pinned. + ...(cedarText.length > 0 && { cedar_policies: cedarText }), + // Registry (#246): the resolved runtime bundle the agent's loaders apply + // (MCP servers merged into .mcp.json; cedar/skills applied downstream). + ...(resolvedAssets.length > 0 && { + resolved_assets: resolvedAssets.map((a) => ({ + kind: a.kind, namespace: a.namespace, name: a.name, version: a.version, runtime: a.runtime, + })), + }), // The agent's PreToolUse hook uses this to compute the maxLifetime // ceiling on per-gate human-in-the-loop approval timeouts. // 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 4fe37aebf..b766d0244 100644 --- a/cdk/src/handlers/shared/repo-config.ts +++ b/cdk/src/handlers/shared/repo-config.ts @@ -60,6 +60,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[]; } /** @@ -88,6 +100,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 = makeDocClient(); diff --git a/cdk/src/stacks/agent.ts b/cdk/src/stacks/agent.ts index a4edd6473..73b889546 100644 --- a/cdk/src/stacks/agent.ts +++ b/cdk/src/stacks/agent.ts @@ -192,6 +192,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`, @@ -933,6 +950,7 @@ export class AgentStack extends Stack { guardrailId: inputGuardrail.guardrailId, guardrailVersion: inputGuardrail.guardrailVersion, attachmentsBucket: attachmentsBucket.bucket, + agentRegistryId: agentRegistry.registryId, // 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.'); + }); +}); From 6efa4877aa30978159fa9c19cbafe05925b92e2c Mon Sep 17 00:00:00 2001 From: bgagent Date: Wed, 5 Aug 2026 11:32:20 -0400 Subject: [PATCH 2/3] =?UTF-8?q?fix(registry):=20address=20review=20on=20th?= =?UTF-8?q?e=20integration=20PR=20=E2=80=94=20validation,=20fail-closed=20?= =?UTF-8?q?load,=20cap=20(#246)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Blueprint / orchestrator: - Validate each typed Blueprint field's ref kind at synth (reject a skill ref under assets.mcpServers, etc.) so a field typo can't silently activate a different asset class. - REMOVE asset columns that go empty on update, so a redeploy that cleared the last mcp_server/cedar_policy_module/skill actually detaches the stale refs. - Persist deprecation warnings: stamp them on resolved_assets and emit a durable registry_asset_warning TaskEvent (was a Lambda log only). Agent loader: - Use an injective MCP server key (drop hyphen->underscore collapse) so acme/foo-bar and acme/foo_bar don't clobber each other. - Normalize the MCP runtime transport -> the SDK's discriminant type key when writing .mcp.json, so a published server the docs describe is actually loaded. - Option C fail-closed: raise RegistryAssetLoadError on infrastructure failures (missing repo_dir, .mcp.json write error) so the task fails rather than running with a pinned-but-absent asset; warn+skip degraded-but-safe cases; return the loaded keys. Policy: - Count registry cedar_policy_module bytes (legacy extra_policies path) in the 64 KiB aggregate cap so a large registry policy can't bypass the bound. Docs: document the opt-in forkBlueprintRepo E2E hook in REGISTRY.md (+ mirror). --- agent/src/pipeline.py | 20 +- agent/src/policy.py | 23 ++- agent/src/registry/loader.py | 131 ++++++++++--- agent/tests/test_entrypoint.py | 23 +++ agent/tests/test_pipeline.py | 81 ++++++++ agent/tests/test_policy_three_outcome.py | 13 ++ agent/tests/test_registry_loader.py | 179 +++++++++++++++--- cdk/src/constructs/blueprint.ts | 52 ++++- cdk/src/handlers/shared/orchestrator.ts | 46 ++++- cdk/test/constructs/blueprint.test.ts | 41 ++++ cdk/test/handlers/orchestrate-task.test.ts | 128 +++++++++++++ docs/design/REGISTRY.md | 19 ++ .../src/content/docs/architecture/Registry.md | 19 ++ 13 files changed, 709 insertions(+), 66 deletions(-) diff --git a/agent/src/pipeline.py b/agent/src/pipeline.py index 5cb5687f9..11413b6ce 100644 --- a/agent/src/pipeline.py +++ b/agent/src/pipeline.py @@ -1155,10 +1155,28 @@ def _on_trace_truncated(max_bytes: int, first_dropped: int) -> None: # Registry assets (#246): merge resolved mcp_server configs into # .mcp.json alongside the channel MCP entry, before the project scan. + # Fail-closed (#246 Option C): apply_resolved_assets raises + # RegistryAssetLoadError on an infrastructure failure (missing + # repo_dir / .mcp.json write error) — we let it propagate so the task + # fails rather than running with a pinned-but-absent asset while the + # audit record claims it was loaded. Degraded-but-safe cases (empty + # runtime) are warn+skip inside the loader. if config.resolved_assets: from registry.loader import apply_resolved_assets - apply_resolved_assets(setup.repo_dir, config.resolved_assets) + loaded_mcp_keys = apply_resolved_assets(setup.repo_dir, config.resolved_assets) + log("TASK", f"Registry: applied {len(loaded_mcp_keys)} mcp_server asset(s)") + # ADR-016 ENFORCEMENT (re-apply after the merge): the registry + # merge writes servers into .mcp.json AFTER the strip above, so a + # registry-published Linear server would otherwise slip back in and + # run under bypassPermissions. Re-strip so the enforcement covers + # registry-sourced entries too, not just repo-committed ones. + if strip_linear_mcp_servers(setup.repo_dir): + log( + "WARN", + "Registry: stripped a Linear MCP server introduced by a resolved " + "asset (ADR-016 — the agent must have no Linear tools)", + ) # Download attachments from S3 (version-pinned, integrity-verified) prepared_attachments: list = [] diff --git a/agent/src/policy.py b/agent/src/policy.py index 3b399a568..a18aa3842 100644 --- a/agent/src/policy.py +++ b/agent/src/policy.py @@ -879,13 +879,26 @@ def __init__( if legacy_extra: soft_text = soft_text + "\n" + "\n".join(legacy_extra) - # 64 KB cap on combined blueprint text (finding #12). Built-ins do - # not count against the cap — they are trusted platform content. - blueprint_text = "".join(filter(None, [blueprint_hard_policies, blueprint_soft_policies])) - if len(blueprint_text.encode("utf-8")) > POLICIES_MAX_BYTES: + # 64 KB cap on combined operator-supplied policy text (finding #12). + # Built-ins do not count — they are trusted platform content. Registry + # cedar_policy_module assets arrive via the legacy ``extra_policies`` + # path, so they MUST be counted here too; otherwise a large registry + # policy bypasses the cap entirely (#246 review). Count the raw operator + # text (pre-synthetic-wrapper) so the bound reflects authored bytes. + operator_text = "".join( + filter( + None, + [ + blueprint_hard_policies, + blueprint_soft_policies, + *(extra_policies or []), + ], + ) + ) + if len(operator_text.encode("utf-8")) > POLICIES_MAX_BYTES: raise ValueError( f"cedar_policies exceeds {POLICIES_MAX_BYTES // 1024} KB cap " - f"({len(blueprint_text.encode('utf-8'))} bytes)" + f"({len(operator_text.encode('utf-8'))} bytes)" ) # Parse + validate annotations on each tier. diff --git a/agent/src/registry/loader.py b/agent/src/registry/loader.py index 284b011ab..4e9f970ef 100644 --- a/agent/src/registry/loader.py +++ b/agent/src/registry/loader.py @@ -32,7 +32,54 @@ def _server_key(asset: dict[str, Any]) -> str: namespace = asset.get("namespace", "") name = asset.get("name", "") - return f"{namespace}__{name}".replace("-", "_") + # Do NOT normalize hyphens to underscores: ``acme/foo-bar`` and + # ``acme/foo_bar`` are distinct registry assets, and collapsing both to + # ``acme__foo_bar`` would silently drop one server (last write wins), so the + # loaded tool surface would diverge from the resolved audit bundle (#246). + # MCP config keys allow hyphens, so the raw components are already a safe, + # injective key. + return f"{namespace}__{name}" + + +def _to_mcp_config(runtime: dict[str, Any], server_key: str) -> dict[str, Any]: + """Normalize a registry mcp_server runtime payload into the ``.mcp.json`` + entry shape the Claude Agent SDK actually consumes. + + The registry contract names the discriminant ``transport`` (``http`` / ``sse`` + / ``stdio``), but the SDK's ``McpServerConfig`` (and the existing + ``channel_mcp`` entries) use the key ``type``. Writing ``transport`` + unchanged produces an entry the agent does not recognize, so a published + server following the documented contract would silently fail to load (#246). + Map ``transport`` → ``type`` and pass the rest through untouched. + + Fail-closed: a structurally invalid payload (http/sse without ``url``, stdio + without ``command``, or an unknown transport) raises + :class:`RegistryAssetLoadError`. Writing a broken ``.mcp.json`` entry would + let the task run with the pinned tool surface silently missing while the + audit bundle claims the asset loaded — exactly the fail-open the resolve-side + validation also guards against (#246 review). + """ + transport = runtime.get("transport") or runtime.get("type") + if transport in ("http", "sse"): + if not runtime.get("url"): + raise RegistryAssetLoadError( + f"{server_key}: {transport} mcp_server runtime is missing 'url'" + ) + elif transport == "stdio": + if not runtime.get("command"): + raise RegistryAssetLoadError( + f"{server_key}: stdio mcp_server runtime is missing 'command'" + ) + else: + raise RegistryAssetLoadError( + f"{server_key}: unknown mcp_server transport {transport!r} " + f"(expected http, sse, or stdio)" + ) + if "transport" not in runtime: + return runtime # already in SDK shape + mapped = {k: v for k, v in runtime.items() if k != "transport"} + mapped["type"] = runtime["transport"] + return mapped def _read_existing_mcp_config(path: str) -> dict[str, Any]: @@ -54,19 +101,42 @@ def _read_existing_mcp_config(path: str) -> dict[str, Any]: return {} -def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> int: +class RegistryAssetLoadError(RuntimeError): + """A resolved asset could not be applied due to an *infrastructure* failure + (the asset resolved fine, but writing it to disk failed). Raised so the task + fails-closed rather than running with an audit record claiming an asset that + was never actually loaded (#246 Option C). Contrast with *degraded-but-safe* + conditions (empty runtime, malformed existing config), which warn + skip.""" + + +def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> list[str]: """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. + Returns the list of server keys actually written. Empty when there are no + mcp_server assets. + + Fail-closed on any condition that would leave a pinned asset unloaded while + the audit bundle claims it loaded (raises :class:`RegistryAssetLoadError`): + * ``repo_dir`` missing / not a directory — the asset resolved but there's + nowhere to write it. + * ``.mcp.json`` write error (OSError). + * an empty / non-dict runtime payload for a pinned asset. + * a structurally invalid connection config (see :func:`_to_mcp_config`). + + A pinned asset is one the operator explicitly referenced in the Blueprint, so + "load it or fail the task" keeps the stamped ``resolved_assets`` audit record + accurate by construction — a warn-and-skip here would let the record claim an + asset the agent never actually loaded (#246 review, Option C). """ mcp_assets = [a for a in resolved_assets if a.get("kind") == _MCP_KIND] if not mcp_assets: - return 0 + return [] 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 + raise RegistryAssetLoadError( + f"cannot apply {len(mcp_assets)} resolved mcp_server asset(s): " + f"repo_dir missing or not a directory: {repo_dir!r}" + ) mcp_path = os.path.join(repo_dir, ".mcp.json") config = _read_existing_mcp_config(mcp_path) @@ -74,17 +144,19 @@ def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> in if not isinstance(servers, dict): servers = {} - written = 0 + written: list[str] = [] for asset in mcp_assets: + key = _server_key(asset) 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 + # Fail closed: a pinned asset with no runtime cannot be honored, and + # skipping it would make the stamped audit bundle lie about what ran. + raise RegistryAssetLoadError(f"{key}: resolved mcp_server has an empty runtime payload") + servers[key] = _to_mcp_config(runtime, key) + written.append(key) - if written == 0: - return 0 + if not written: + return [] config["mcpServers"] = servers try: @@ -92,10 +164,9 @@ def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> in 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 + raise RegistryAssetLoadError(f"failed to write {mcp_path}: {e}") from e - log("TASK", f"Registry: merged {written} MCP server(s) into {mcp_path}") + log("TASK", f"Registry: merged {len(written)} MCP server(s) into {mcp_path}") return written @@ -117,32 +188,36 @@ def build_skill_prompt_fragment(resolved_assets: list[dict[str, Any]]) -> str: 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', '')}" + runtime = asset.get("runtime") + fragment = runtime.get("prompt_fragment") if isinstance(runtime, dict) else None + if not isinstance(runtime, dict) or not isinstance(fragment, str) or not fragment.strip(): + # Fail closed: a pinned skill whose fragment is missing/empty would be + # silently dropped from the prompt while still stamped as loaded in the + # audit bundle — surface it instead (#246 review, Option C). + raise RegistryAssetLoadError(f"{name}: resolved skill has no usable 'prompt_fragment'") 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: +def apply_resolved_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> list[str]: """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. + + Returns the list of mcp_server keys actually written (for the caller to + reconcile against the stamped audit bundle). Propagates + :class:`RegistryAssetLoadError` on an infrastructure failure so the pipeline + fails the task rather than running with a resolved asset silently missing. """ if not resolved_assets: - return - apply_mcp_assets(repo_dir, resolved_assets) + return [] + return apply_mcp_assets(repo_dir, resolved_assets) diff --git a/agent/tests/test_entrypoint.py b/agent/tests/test_entrypoint.py index 96afdb3b2..2ec7a84f9 100644 --- a/agent/tests/test_entrypoint.py +++ b/agent/tests/test_entrypoint.py @@ -372,6 +372,29 @@ def test_overrides_appended(self): assert "Always use tabs" in result assert "Additional instructions" in result + def test_resolved_skill_fragment_appended_to_system_prompt(self): + # #246: a resolved skill's prompt_fragment must reach the system prompt. + config = TaskConfig( + repo_url="o/r", + task_id="t1", + max_turns=10, + github_token="ghp_test", + aws_region="us-east-1", + resolved_assets=[ + { + "kind": "skill", + "namespace": "acme", + "name": "readme-helper", + "version": "1.0.0", + "runtime": {"prompt_fragment": "Add an ABCA-REVIEWED marker."}, + } + ], + ) + setup = RepoSetup(repo_dir="/workspace/t1", branch="b", default_branch="main", notes=[]) + result = _build_system_prompt(config, setup, None, "") + assert "## Skills" in result + assert "Add an ABCA-REVIEWED marker." in result + # --------------------------------------------------------------------------- # build_config — workflow resolution diff --git a/agent/tests/test_pipeline.py b/agent/tests/test_pipeline.py index 4fcff8797..e6b26ea6d 100644 --- a/agent/tests/test_pipeline.py +++ b/agent/tests/test_pipeline.py @@ -159,6 +159,87 @@ async def fake_run_agent(_prompt, _system_prompt, config, cwd=None, trajectory=N assert captured_config is not None assert captured_config.cedar_policies == [] + @patch("runner.run_agent") + @patch("pipeline.build_system_prompt") + @patch("pipeline.discover_project_config") + @patch("repo.setup_repo") + @patch("pipeline.task_span") + @patch("pipeline.task_state") + def test_malformed_registry_asset_fails_the_task_closed( + self, + mock_task_state, + mock_task_span, + mock_setup_repo, + _mock_discover, + _mock_build_prompt, + mock_run_agent, + monkeypatch, + tmp_path, + ): + """#246 fail-closed: a resolved mcp_server whose runtime is structurally + invalid must fail the task (write_terminal FAILED) and re-raise, never run + the agent with the pinned asset silently missing.""" + monkeypatch.setenv("GITHUB_TOKEN", "ghp_test") + monkeypatch.setenv("AWS_REGION", "us-east-1") + + # A real repo_dir so the loader reaches the transport-validation branch + # (the failure we want is the invalid payload, not a missing dir). + mock_setup_repo.return_value = RepoSetup( + repo_dir=str(tmp_path), + branch="bgagent/test/branch", + build_before=True, + ) + + agent_ran = False + + async def fake_run_agent(_prompt, _system_prompt, config, cwd=None, trajectory=None): + nonlocal agent_ran + agent_ran = True + return AgentResult(status="success", turns=1, cost_usd=0.01, num_turns=1) + + mock_run_agent.side_effect = fake_run_agent + + mock_span = MagicMock() + mock_span.__enter__ = MagicMock(return_value=mock_span) + mock_span.__exit__ = MagicMock(return_value=False) + mock_task_span.return_value = mock_span + mock_task_state.get_task.return_value = None + + with ( + patch("pipeline.configure_channel_mcp"), + patch("pipeline.strip_linear_mcp_servers", return_value=0), + patch("pipeline.get_disk_usage", return_value=0), + patch("pipeline.print_metrics"), + ): + from pipeline import run_task + + # http transport with no url → RegistryAssetLoadError inside the loader. + bad_asset = { + "kind": "mcp_server", + "namespace": "acme", + "name": "pdf-tools", + "version": "1.0.0", + "runtime": {"transport": "http"}, + } + with pytest.raises(Exception): # noqa: B017 — re-raised after FAILED write + run_task( + repo_url="owner/repo", + task_description="fix bug", + github_token="ghp_test", + aws_region="us-east-1", + task_id="test-id", + resolved_assets=[bad_asset], + ) + + # The task was marked FAILED and the agent never ran with a missing asset. + assert agent_ran is False + failed_writes = [ + c for c in mock_task_state.write_terminal.call_args_list if c.args[1] == "FAILED" + ] + assert failed_writes, ( + "expected a write_terminal(..., 'FAILED', ...) on the fail-closed path" + ) + @patch("runner.run_agent") @patch("pipeline.build_system_prompt") @patch("pipeline.discover_project_config") diff --git a/agent/tests/test_policy_three_outcome.py b/agent/tests/test_policy_three_outcome.py index 0adc8988e..244977d4e 100644 --- a/agent/tests/test_policy_three_outcome.py +++ b/agent/tests/test_policy_three_outcome.py @@ -715,6 +715,19 @@ def test_blueprint_64kb_cap_rejected(self): blueprint_soft_policies=big, ) + def test_registry_extra_policies_counted_in_64kb_cap(self): + # #246 review: registry cedar_policy_module assets arrive via the legacy + # extra_policies path, which previously bypassed the cap entirely. An + # oversized registry policy must be rejected just like blueprint text. + big = 'forbid (principal, action, resource) when { context.x like "*aaaaaaaaaa*" };' * 1000 + assert len(big) > POLICIES_MAX_BYTES + with pytest.raises(ValueError, match="64 KB cap"): + PolicyEngine( + task_type="new_task", + repo="owner/repo", + extra_policies=[big], + ) + def test_blueprint_soft_rule_missing_rule_id_rejected(self): bad = '@tier("soft") forbid (principal, action, resource) when { context.x like "*foo*" };' with pytest.raises(ValueError, match="missing @rule_id"): diff --git a/agent/tests/test_registry_loader.py b/agent/tests/test_registry_loader.py index ae351efee..ed06ad831 100644 --- a/agent/tests/test_registry_loader.py +++ b/agent/tests/test_registry_loader.py @@ -4,7 +4,10 @@ import json +import pytest + from registry.loader import ( + RegistryAssetLoadError, apply_mcp_assets, apply_resolved_assets, build_skill_prompt_fragment, @@ -29,19 +32,43 @@ def _mcp_asset(namespace: str, name: str, version: str, runtime: dict) -> dict: 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 + asset = _mcp_asset("acme", "pdf-tools", "1.0.0", runtime) + written = apply_mcp_assets(str(tmp_path), [asset]) + assert written == ["acme__pdf-tools"] merged = _read_mcp(tmp_path) - assert merged["mcpServers"]["acme__pdf_tools"] == runtime + # Hyphens are preserved (injective key) — not normalized to underscores. + # `transport` is normalized to the SDK's `type` discriminant key (#246). + assert merged["mcpServers"]["acme__pdf-tools"] == { + "type": "http", + "url": "https://mcp.example.com/sse", + } + + def test_normalizes_transport_to_type(self, tmp_path): + # A publisher following the documented `transport` contract must produce + # a `.mcp.json` entry the SDK recognizes (discriminant key `type`), with + # no leftover `transport` key and all other fields preserved. + runtime = { + "transport": "sse", + "url": "https://x/sse", + "headers": {"Authorization": "Bearer t"}, + "tool_prefix": "mcp__x__", + } + apply_mcp_assets(str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", runtime)]) + entry = _read_mcp(tmp_path)["mcpServers"]["acme__x"] + assert entry["type"] == "sse" + assert "transport" not in entry + assert entry["url"] == "https://x/sse" + assert entry["headers"] == {"Authorization": "Bearer t"} + assert entry["tool_prefix"] == "mcp__x__" 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( + written = apply_mcp_assets( str(tmp_path), [_mcp_asset("acme", "weather", "2.1.0", {"transport": "sse", "url": "https://w"})], ) - assert n == 1 + assert written == ["acme__weather"] merged = _read_mcp(tmp_path) assert merged["mcpServers"]["other"]["command"] == "/usr/bin/x" assert "acme__weather" in merged["mcpServers"] @@ -51,11 +78,27 @@ def test_merges_multiple_servers(self, tmp_path): _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 + written = apply_mcp_assets(str(tmp_path), assets) + assert set(written) == {"acme__a", "acme__b"} merged = _read_mcp(tmp_path) assert set(merged["mcpServers"]) == {"acme__a", "acme__b"} + def test_hyphen_and_underscore_names_do_not_collide(self, tmp_path): + # foo-bar and foo_bar are distinct assets; the key must not collapse them + # to one entry (last-write-wins would drop a resolved server) (#246). + assets = [ + _mcp_asset("acme", "foo-bar", "1.0.0", {"transport": "http", "url": "https://dash"}), + _mcp_asset( + "acme", "foo_bar", "1.0.0", {"transport": "http", "url": "https://underscore"} + ), + ] + written = apply_mcp_assets(str(tmp_path), assets) + assert set(written) == {"acme__foo-bar", "acme__foo_bar"} + merged = _read_mcp(tmp_path) + assert set(merged["mcpServers"]) == {"acme__foo-bar", "acme__foo_bar"} + assert merged["mcpServers"]["acme__foo-bar"]["url"] == "https://dash" + assert merged["mcpServers"]["acme__foo_bar"]["url"] == "https://underscore" + def test_ignores_non_mcp_kinds(self, tmp_path): assets = [ { @@ -66,26 +109,72 @@ def test_ignores_non_mcp_kinds(self, tmp_path): "runtime": {"cedar_text": "permit(...);"}, }, ] - n = apply_mcp_assets(str(tmp_path), assets) - assert n == 0 + written = apply_mcp_assets(str(tmp_path), assets) + assert written == [] 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 + def test_empty_runtime_raises(self, tmp_path): + # Fail-closed (#246 review): a pinned asset with no runtime cannot be + # honored; skipping it would make the stamped audit bundle lie. + with pytest.raises(RegistryAssetLoadError, match="empty runtime payload"): + apply_mcp_assets(str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", {})]) assert not (tmp_path / ".mcp.json").exists() - def test_no_op_on_missing_repo_dir(self): + def test_http_without_url_raises(self, tmp_path): + with pytest.raises(RegistryAssetLoadError, match="missing 'url'"): + apply_mcp_assets( + str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", {"transport": "http"})] + ) + + def test_stdio_without_command_raises(self, tmp_path): + with pytest.raises(RegistryAssetLoadError, match="missing 'command'"): + apply_mcp_assets( + str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", {"transport": "stdio"})] + ) + + def test_unknown_transport_raises(self, tmp_path): + with pytest.raises(RegistryAssetLoadError, match="unknown mcp_server transport"): + apply_mcp_assets( + str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", {"transport": "grpc", "url": "u"})] + ) + + def test_stdio_with_command_loads(self, tmp_path): + written = apply_mcp_assets( + str(tmp_path), + [_mcp_asset("acme", "x", "1.0.0", {"transport": "stdio", "command": "run-me"})], + ) + assert written == ["acme__x"] + assert _read_mcp(tmp_path)["mcpServers"]["acme__x"] == { + "type": "stdio", + "command": "run-me", + } + + def test_missing_repo_dir_raises(self): + # Infrastructure failure (#246 Option C): the asset resolved but there's + # nowhere to write it — fail-closed so the audit can't claim it loaded. + asset = _mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "u"}) + with pytest.raises(RegistryAssetLoadError, match="repo_dir missing"): + apply_mcp_assets("/nonexistent/dir", [asset]) + + def test_write_error_raises(self, tmp_path, monkeypatch): + # Infrastructure failure: .mcp.json write fails → fail-closed. asset = _mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "u"}) - n = apply_mcp_assets("/nonexistent/dir", [asset]) - assert n == 0 + + def _boom(*_a, **_k): + raise OSError("disk full") + + monkeypatch.setattr("builtins.open", _boom) + with pytest.raises(RegistryAssetLoadError, match="failed to write"): + apply_mcp_assets(str(tmp_path), [asset]) def test_malformed_existing_treated_as_absent(self, tmp_path): + # Degraded-but-safe: a corrupt existing .mcp.json is replaced, not fatal. (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 + written = apply_mcp_assets(str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", runtime)]) + assert written == ["acme__x"] + # Written in SDK shape (transport → type). + assert _read_mcp(tmp_path)["mcpServers"]["acme__x"] == {"type": "http", "url": "https://x"} def _skill_asset(namespace: str, name: str, runtime: dict) -> dict: @@ -126,10 +215,16 @@ def test_concatenates_multiple_in_order(self): ) assert out.index("First.") < out.index("Second.") - def test_skips_blank_or_invalid_runtime(self): + def test_blank_fragment_raises(self): + # Fail-closed (#246 review): a pinned skill whose fragment is missing/blank + # would otherwise be silently dropped while stamped as loaded. blank = _skill_asset("acme", "a", {"prompt_fragment": " "}) - assert build_skill_prompt_fragment([blank]) == "" - assert build_skill_prompt_fragment([_skill_asset("acme", "a", {})]) == "" + with pytest.raises(RegistryAssetLoadError, match="no usable 'prompt_fragment'"): + build_skill_prompt_fragment([blank]) + + def test_missing_runtime_raises(self): + with pytest.raises(RegistryAssetLoadError, match="no usable 'prompt_fragment'"): + build_skill_prompt_fragment([_skill_asset("acme", "a", {})]) class TestApplyResolvedAssets: @@ -138,14 +233,54 @@ def test_empty_is_noop(self, tmp_path): assert not (tmp_path / ".mcp.json").exists() def test_dispatches_mcp(self, tmp_path): - apply_resolved_assets( + written = apply_resolved_assets( str(tmp_path), [_mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "https://x"})], ) + assert written == ["acme__x"] assert (tmp_path / ".mcp.json").exists() + def test_propagates_infra_failure(self, tmp_path): + # Fail-closed (#246 Option C): an mcp_server that resolved but can't be + # written must raise so the pipeline fails the task. + with pytest.raises(RegistryAssetLoadError): + apply_resolved_assets( + "/nonexistent/dir", + [_mcp_asset("acme", "x", "1.0.0", {"transport": "http", "url": "u"})], + ) + 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() + + +class TestAdr016LinearReStrip: + """A registry-published Linear MCP server merged into .mcp.json must be + scrubbed by strip_linear_mcp_servers (ADR-016), which the pipeline now runs + AFTER the registry merge. Guards the bypass where a registry asset could + re-introduce Linear tools under bypassPermissions (#246 review).""" + + def test_registry_linear_server_is_stripped_after_merge(self, tmp_path): + from channel_mcp import strip_linear_mcp_servers + + # A registry asset that (maliciously or accidentally) provides Linear. + apply_resolved_assets( + str(tmp_path), + [ + _mcp_asset( + "evil", + "linear", + "1.0.0", + {"transport": "http", "url": "https://mcp.linear.app/sse"}, + ), + _mcp_asset("acme", "pdf", "1.0.0", {"transport": "http", "url": "https://pdf"}), + ], + ) + # The pipeline runs this immediately after the merge. + removed = strip_linear_mcp_servers(str(tmp_path)) + assert removed == 1 + servers = _read_mcp(tmp_path)["mcpServers"] + assert "evil__linear" not in servers # Linear scrubbed + assert "acme__pdf" in servers # benign server survives diff --git a/cdk/src/constructs/blueprint.ts b/cdk/src/constructs/blueprint.ts index 751d52c3c..dbb2c8c3d 100644 --- a/cdk/src/constructs/blueprint.ts +++ b/cdk/src/constructs/blueprint.ts @@ -258,9 +258,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)); + this.node.addValidation(new RegistryRefValidation('assets.mcpServers', this.mcpServerRefs, 'mcp_server')); + this.node.addValidation(new RegistryRefValidation('assets.cedarPolicyModules', this.cedarPolicyModuleRefs, 'cedar_policy_module')); + this.node.addValidation(new RegistryRefValidation('assets.skills', this.skillRefs, 'skill')); const now = new Date().toISOString(); @@ -335,11 +335,12 @@ export class Blueprint extends Construct { parameters: { TableName: props.repoTable.tableName, Key: { repo: { S: props.repo } }, - UpdateExpression: `SET #status = :active, #updated = :now${this.buildUpdateFields(props)}`, + UpdateExpression: `SET #status = :active, #updated = :now${this.buildUpdateFields(props)}${this.buildRemoveClause()}`, ExpressionAttributeNames: { '#status': 'status', '#updated': 'updated_at', ...this.buildExpressionNames(props), + ...this.buildRemoveNames(), }, ExpressionAttributeValues: { ':active': { S: 'active' }, @@ -438,6 +439,29 @@ export class Blueprint extends Construct { if (this.skillRefs.length > 0) values[':skills'] = { L: this.skillRefs.map(r => ({ S: r })) }; return values; } + + /** Registry asset fields that are now empty must be REMOVEd on update, not + * just omitted from SET — otherwise a redeploy that cleared the last + * mcp_server/cedar_policy_module/skill leaves the stale DDB refs active and + * operators can't detach a pinned asset through the Blueprint API (#246). */ + private emptyAssetFields(): string[] { + const empty: string[] = []; + if (this.mcpServerRefs.length === 0) empty.push('mcp_servers'); + if (this.cedarPolicyModuleRefs.length === 0) empty.push('cedar_policy_modules'); + if (this.skillRefs.length === 0) empty.push('skills'); + return empty; + } + + private buildRemoveClause(): string { + const empty = this.emptyAssetFields(); + return empty.length > 0 ? ` REMOVE ${empty.map(f => `#${f}`).join(', ')}` : ''; + } + + private buildRemoveNames(): Record { + const names: Record = {}; + for (const f of this.emptyAssetFields()) names[`#${f}`] = f; + return names; + } } /** @@ -502,9 +526,20 @@ class ApprovalGateCapValidation implements IValidation { * 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. + * + * Also enforces that the ref's kind matches the field it was pinned under + * (``expectedKind``). Each typed Blueprint field stores into a distinct DDB + * column, and the orchestrator dispatches by the ref's embedded kind — so a + * ``skill`` ref placed under ``assets.mcpServers`` would otherwise deploy and + * then silently activate skill behavior from an "MCP" column. Reject the + * mismatch at synth instead. */ class RegistryRefValidation implements IValidation { - constructor(private readonly field: string, private readonly refs: readonly string[]) {} + constructor( + private readonly field: string, + private readonly refs: readonly string[], + private readonly expectedKind: string, + ) {} public validate(): string[] { const errors: string[] = []; @@ -512,6 +547,13 @@ class RegistryRefValidation implements IValidation { const result = parseRef(ref); if (!result.ok) { errors.push(`Invalid ${this.field} ref '${ref}': ${result.reason} — ${result.message}`); + continue; + } + if (result.ref.kind !== this.expectedKind) { + errors.push( + `Wrong kind for ${this.field} ref '${ref}': expected a '${this.expectedKind}' ref ` + + `but got '${result.ref.kind}'.`, + ); } } return errors; diff --git a/cdk/src/handlers/shared/orchestrator.ts b/cdk/src/handlers/shared/orchestrator.ts index 77bce1786..36ab2d74d 100644 --- a/cdk/src/handlers/shared/orchestrator.ts +++ b/cdk/src/handlers/shared/orchestrator.ts @@ -800,22 +800,58 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B 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 })), + // Persist warnings (e.g. ["DEPRECATED"]) alongside the audit triple so a + // user inspecting the task record can see a deprecated asset ran — ADR-022 + // sub-decision 4 promises this, and a Lambda log alone isn't durable (#246). + ':ra': resolvedAssets.map((a) => ({ + kind: a.kind, + id: `${a.namespace}/${a.name}`, + version: a.version, + ...(a.warnings.length > 0 && { warnings: [...a.warnings] }), + })), ':now': new Date().toISOString(), }, })); + + // Emit a durable TaskEvent per warned asset (deprecation is the main case), + // so the warning surfaces in the task's event stream, not just Lambda logs. + for (const a of resolvedAssets) { + if (a.warnings.length > 0) { + await emitTaskEvent(task.task_id, 'registry_asset_warning', { + kind: a.kind, + id: `${a.namespace}/${a.name}`, + version: a.version, + warnings: [...a.warnings], + }, correlation); + } + } } // 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. + // + // Fail-closed: a pinned module whose cedar_text is empty/whitespace must fail + // the task, not be silently dropped — a dropped policy is usually a *deny* rule, + // so silently omitting it would WIDEN what the agent may do while the audit + // record still claims the module was applied (#246 review). + const resolvedCedar = resolvedAssets + .filter((a) => a.kind === 'cedar_policy_module') + .map((a) => { + const text = (a.runtime as { cedar_text?: string }).cedar_text; + if (typeof text !== 'string' || text.trim().length === 0) { + throw new RegistryResolutionError( + 'REMOVED', + `registry://cedar_policy_module/${a.namespace}/${a.name}@${a.version}`, + `resolved cedar_policy_module ${a.namespace}/${a.name}@${a.version} has empty cedar_text`, + ); + } + return text; + }); 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), + ...resolvedCedar, ]; const payload: Record = { diff --git a/cdk/test/constructs/blueprint.test.ts b/cdk/test/constructs/blueprint.test.ts index b1ac03eab..482e691cf 100644 --- a/cdk/test/constructs/blueprint.test.ts +++ b/cdk/test/constructs/blueprint.test.ts @@ -337,6 +337,29 @@ describe('Blueprint construct', () => { expect(serialized).toContain('#mcp_servers'); expect(serialized).toContain('#cedar_policy_modules'); expect(serialized).toContain('#skills'); + // All three populated → nothing to REMOVE. + expect(serialized).not.toContain('REMOVE'); + }); + + test('onUpdate REMOVEs asset columns that are now empty (detach on redeploy)', () => { + // Only mcpServers pinned: cedar_policy_modules + skills must be REMOVEd so a + // redeploy that cleared them detaches the stale DDB refs (#246). + const { template } = createStack({ + assets: { mcpServers: ['registry://mcp_server/acme/pdf-tools@^1.4.1'] }, + }); + const serialized = getUpdateJoinParts(template).join(''); + // mcp_servers is SET (populated); the other two are REMOVEd. Assert on the + // exact REMOVE clause so the ExpressionAttributeNames block (which maps all + // three names) doesn't confuse the check. + expect(serialized).toContain('#mcp_servers = :mcp_servers'); + expect(serialized).toContain('REMOVE #cedar_policy_modules, #skills'); + expect(serialized).not.toContain('REMOVE #mcp_servers'); + }); + + test('onUpdate REMOVEs all three asset columns when none are pinned', () => { + const { template } = createStack(); + const serialized = getUpdateJoinParts(template).join(''); + expect(serialized).toContain('REMOVE #mcp_servers, #cedar_policy_modules, #skills'); }); test('rejects a floating asset ref at synth', () => { @@ -367,6 +390,24 @@ describe('Blueprint construct', () => { expect(() => Template.fromStack(stack)).toThrow(/Invalid assets.skills ref.*INVALID_CONSTRAINT/); }); + test('rejects a ref whose kind does not match its field 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, + // A well-formed skill ref, but pinned under mcpServers — must be rejected + // so a field typo can't silently activate a different asset class. + assets: { mcpServers: ['registry://skill/acme/research@1.0.0'] }, + }); + expect(() => Template.fromStack(stack)).toThrow( + /Wrong kind for assets.mcpServers ref.*expected a 'mcp_server' ref but got 'skill'/, + ); + }); + // --- Chunk 7b: security.approvalGateCap --------------------------------- test('exposes approvalGateCap as public property when configured', () => { diff --git a/cdk/test/handlers/orchestrate-task.test.ts b/cdk/test/handlers/orchestrate-task.test.ts index 42247705c..ee3e81676 100644 --- a/cdk/test/handlers/orchestrate-task.test.ts +++ b/cdk/test/handlers/orchestrate-task.test.ts @@ -55,6 +55,14 @@ jest.mock('../../src/handlers/shared/repo-config', () => ({ checkRepoOnboarded: jest.fn(), })); +// Registry client (#246): the orchestrator resolves Blueprint registry:// refs +// through this factory. Mock it so hydrateAndTransition tests can drive the +// resolve → stamp → payload path without the AgentCore SDK. +const mockRegistryResolve = jest.fn(); +jest.mock('../../src/handlers/shared/registry/factory', () => ({ + makeRegistryClient: jest.fn(() => ({ resolve: mockRegistryResolve })), +})); + let ulidCounter = 0; jest.mock('ulid', () => ({ ulid: jest.fn(() => `ULID${ulidCounter++}`) })); @@ -80,6 +88,7 @@ import { queueTask, transitionTask, } from '../../src/handlers/shared/orchestrator'; +import { RegistryResolutionError } from '../../src/handlers/shared/registry/types'; const baseTask = { task_id: 'TASK001', @@ -897,6 +906,125 @@ describe('hydrateAndTransition with blueprint config', () => { }); }); +describe('hydrateAndTransition — registry asset resolution (#246)', () => { + const mockHydratedContext = { + version: 1, + user_prompt: 'Task ID: TASK001\nRepository: org/repo\n\n## Task\n\nFix the bug', + sources: ['task_description'], + token_estimate: 20, + truncated: false, + content_trust: { task_description: 'trusted' }, + }; + + const resolvedAsset = (over: Record) => ({ + kind: 'mcp_server', + namespace: 'acme', + name: 'pdf-tools', + version: '1.0.0', + runtime: { transport: 'http', url: 'https://x' }, + warnings: [], + ...over, + }); + + test('stamps resolved_assets on the TaskRecord and threads the bundle into the payload', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + mockRegistryResolve.mockResolvedValueOnce(resolvedAsset({})); + + const payload = await hydrateAndTransition(baseTask as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + mcp_servers: ['registry://mcp_server/acme/pdf-tools@^1.0.0'], + }); + + // Threaded into the agent payload. + expect(payload.resolved_assets).toEqual([ + expect.objectContaining({ kind: 'mcp_server', namespace: 'acme', name: 'pdf-tools', version: '1.0.0' }), + ]); + // Stamped on the TaskRecord via an UpdateCommand carrying resolved_assets. + const stamp = mockDdbSend.mock.calls + .map((c) => c[0]) + .find((cmd: any) => cmd._type === 'Update' && cmd.input?.ExpressionAttributeNames?.['#ra'] === 'resolved_assets'); + expect(stamp).toBeDefined(); + expect(stamp.input.ExpressionAttributeValues[':ra']).toEqual([ + { kind: 'mcp_server', id: 'acme/pdf-tools', version: '1.0.0' }, + ]); + }); + + test('emits a registry_asset_warning TaskEvent for a DEPRECATED asset', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + mockRegistryResolve.mockResolvedValueOnce(resolvedAsset({ warnings: ['DEPRECATED'] })); + + await hydrateAndTransition(baseTask as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + mcp_servers: ['registry://mcp_server/acme/pdf-tools@^1.0.0'], + }); + + // A TaskEvent Put (event_type registry_asset_warning) was written. + const warnPut = mockDdbSend.mock.calls + .map((c) => c[0]) + .find((cmd: any) => cmd._type === 'Put' && JSON.stringify(cmd.input?.Item ?? {}).includes('registry_asset_warning')); + expect(warnPut).toBeDefined(); + // The stamped audit triple keeps the warning too. + const stamp = mockDdbSend.mock.calls + .map((c) => c[0]) + .find((cmd: any) => cmd._type === 'Update' && cmd.input?.ExpressionAttributeNames?.['#ra'] === 'resolved_assets'); + expect(stamp.input.ExpressionAttributeValues[':ra'][0].warnings).toEqual(['DEPRECATED']); + }); + + test('merges resolved cedar_policy_module text after inline blueprint policies', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + mockRegistryResolve.mockResolvedValueOnce(resolvedAsset({ + kind: 'cedar_policy_module', + name: 'guard', + runtime: { cedar_text: 'forbid (principal, action, resource);' }, + })); + + const payload = await hydrateAndTransition(baseTask as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + cedar_policies: ['permit (principal, action, resource);'], + cedar_policy_modules: ['registry://cedar_policy_module/acme/guard@^1.0.0'], + }); + + expect(payload.cedar_policies).toEqual([ + 'permit (principal, action, resource);', + 'forbid (principal, action, resource);', + ]); + }); + + test('fails closed when a pinned cedar_policy_module resolves to empty cedar_text', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + mockRegistryResolve.mockResolvedValueOnce(resolvedAsset({ + kind: 'cedar_policy_module', + name: 'guard', + runtime: { cedar_text: ' ' }, + })); + + await expect(hydrateAndTransition(baseTask as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + cedar_policy_modules: ['registry://cedar_policy_module/acme/guard@^1.0.0'], + })).rejects.toThrow(/empty cedar_text/); + }); + + test('fails closed (propagates) when a registry ref cannot be resolved', async () => { + mockDdbSend.mockResolvedValue({}); + mockHydrateContext.mockResolvedValueOnce(mockHydratedContext); + mockRegistryResolve.mockRejectedValueOnce(new RegistryResolutionError('NO_MATCHING_VERSION', 'r', 'none')); + + await expect(hydrateAndTransition(baseTask as any, { + compute_type: 'agentcore', + runtime_arn: 'arn:test', + mcp_servers: ['registry://mcp_server/acme/pdf-tools@^9.0.0'], + })).rejects.toBeInstanceOf(RegistryResolutionError); + }); +}); + describe('finalizeTask', () => { test('handles already-terminal task', async () => { mockDdbSend diff --git a/docs/design/REGISTRY.md b/docs/design/REGISTRY.md index 9e0ed6824..19646b3b7 100644 --- a/docs/design/REGISTRY.md +++ b/docs/design/REGISTRY.md @@ -186,6 +186,25 @@ The strict grammar is implemented by `parseRef` (TS) and `parse_ref` (Python), k - **Construct tests**: `registry.test.ts` (Provider wiring + IAM). - **E2E (PR 2)**: publish an MCP server → reference from a Blueprint → run a task → assert the agent payload carries the bundle and the `TaskRecord` has `resolved_assets`. +### 12.1 Reproducing the E2E — the `forkBlueprintRepo` demo hook + +The stack ships an **opt-in** deploy hook that onboards one repo with all three MVP asset kinds pinned, so the end-to-end path can be exercised without hand-authoring a Blueprint. It is off by default (no fork is hardcoded for other contributors). Enable it by pointing it at a repo you control: + +```bash +# via CDK context… +cdk deploy --context forkBlueprintRepo=owner/repo +# …or via env var +FORK_BLUEPRINT_REPO=owner/repo cdk deploy +``` + +When set, the stack adds a `Blueprint` for `owner/repo` pinning +`registry://mcp_server/acme/aws-knowledge@^1.0.0`, +`registry://cedar_policy_module/acme/guard@^1.0.0`, and +`registry://skill/acme/readme-helper@^1.0.0`. Those `acme/*` records must be +published to the registry first (they are illustrative, not seeded) — otherwise +task admission fails closed on the unresolved pins. Leave the flag unset for a +normal deploy. + ## 13. Accepted risk Preview API: AgentCore Registry hard-migrates namespaces at GA (~2026-08-06) with breaking API-schema changes. The `RegistryClient` port confines the rework to one adapter file per language; experimental project + no prod data ⇒ acceptable. Swap the provisioning custom resource for native CDK constructs when they ship at GA. diff --git a/docs/src/content/docs/architecture/Registry.md b/docs/src/content/docs/architecture/Registry.md index 849c7c52c..bdca43bdb 100644 --- a/docs/src/content/docs/architecture/Registry.md +++ b/docs/src/content/docs/architecture/Registry.md @@ -190,6 +190,25 @@ The strict grammar is implemented by `parseRef` (TS) and `parse_ref` (Python), k - **Construct tests**: `registry.test.ts` (Provider wiring + IAM). - **E2E (PR 2)**: publish an MCP server → reference from a Blueprint → run a task → assert the agent payload carries the bundle and the `TaskRecord` has `resolved_assets`. +### 12.1 Reproducing the E2E — the `forkBlueprintRepo` demo hook + +The stack ships an **opt-in** deploy hook that onboards one repo with all three MVP asset kinds pinned, so the end-to-end path can be exercised without hand-authoring a Blueprint. It is off by default (no fork is hardcoded for other contributors). Enable it by pointing it at a repo you control: + +```bash +# via CDK context… +cdk deploy --context forkBlueprintRepo=owner/repo +# …or via env var +FORK_BLUEPRINT_REPO=owner/repo cdk deploy +``` + +When set, the stack adds a `Blueprint` for `owner/repo` pinning +`registry://mcp_server/acme/aws-knowledge@^1.0.0`, +`registry://cedar_policy_module/acme/guard@^1.0.0`, and +`registry://skill/acme/readme-helper@^1.0.0`. Those `acme/*` records must be +published to the registry first (they are illustrative, not seeded) — otherwise +task admission fails closed on the unresolved pins. Leave the flag unset for a +normal deploy. + ## 13. Accepted risk Preview API: AgentCore Registry hard-migrates namespaces at GA (~2026-08-06) with breaking API-schema changes. The `RegistryClient` port confines the rework to one adapter file per language; experimental project + no prod data ⇒ acceptable. Swap the provisioning custom resource for native CDK constructs when they ship at GA. From ee4011ff043f9494906ca5229406f195f28c90ef Mon Sep 17 00:00:00 2001 From: bgagent Date: Tue, 11 Aug 2026 22:33:09 -0400 Subject: [PATCH 3/3] fix(registry): block .mcp.json secret exfiltration + scope orchestrator IAM (#246) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third review pass (@scottschreckengaust) on the integration PR: - B4 (P1 security): the resolved runtime may carry secrets (bearer headers, url tokens, --api-key args) and .mcp.json lives in the live git clone, so the post-hook `git add -u` → commit → push could exfiltrate it to the PR when the target repo tracks .mcp.json. apply_mcp_assets now marks the file skip-worktree (intent-to-add first if untracked), blocking both `git add -u` and an explicit add. Regression tests assert a secret-bearing asset cannot be staged (tracked + untracked cases). - nit: scope the orchestrator's registry read grant to registry/{agentRegistryId} (+ record/* suffix) instead of registry/* — mirrors registry-api.ts; nag reason trimmed; test asserts no bare "*"/registry-id wildcard. - nits: align the RegistryAssetLoadError docstring + pipeline comment with the fail-closed behavior (no more stale "warn+skip"); drop the now-dead `if not written` branch. --- agent/src/pipeline.py | 8 +-- agent/src/registry/loader.py | 66 ++++++++++++++++--- agent/tests/test_registry_loader.py | 62 +++++++++++++++++ cdk/src/constructs/task-orchestrator.ts | 12 ++-- cdk/test/constructs/task-orchestrator.test.ts | 24 +++++++ 5 files changed, 155 insertions(+), 17 deletions(-) diff --git a/agent/src/pipeline.py b/agent/src/pipeline.py index 11413b6ce..709301d7c 100644 --- a/agent/src/pipeline.py +++ b/agent/src/pipeline.py @@ -1156,11 +1156,11 @@ def _on_trace_truncated(max_bytes: int, first_dropped: int) -> None: # Registry assets (#246): merge resolved mcp_server configs into # .mcp.json alongside the channel MCP entry, before the project scan. # Fail-closed (#246 Option C): apply_resolved_assets raises - # RegistryAssetLoadError on an infrastructure failure (missing - # repo_dir / .mcp.json write error) — we let it propagate so the task + # RegistryAssetLoadError for any condition that would leave a pinned + # asset unloaded (missing repo_dir, empty/invalid runtime, structurally + # invalid config, or a write error) — we let it propagate so the task # fails rather than running with a pinned-but-absent asset while the - # audit record claims it was loaded. Degraded-but-safe cases (empty - # runtime) are warn+skip inside the loader. + # audit record claims it was loaded. if config.resolved_assets: from registry.loader import apply_resolved_assets diff --git a/agent/src/registry/loader.py b/agent/src/registry/loader.py index 4e9f970ef..93f3d596c 100644 --- a/agent/src/registry/loader.py +++ b/agent/src/registry/loader.py @@ -18,10 +18,14 @@ import json import os +import subprocess from typing import Any from shell import log +# Seconds to allow each git plumbing call when protecting .mcp.json from commit. +_GIT_GUARD_TIMEOUT_S = 15 + # 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. @@ -102,11 +106,56 @@ def _read_existing_mcp_config(path: str) -> dict[str, Any]: class RegistryAssetLoadError(RuntimeError): - """A resolved asset could not be applied due to an *infrastructure* failure - (the asset resolved fine, but writing it to disk failed). Raised so the task - fails-closed rather than running with an audit record claiming an asset that - was never actually loaded (#246 Option C). Contrast with *degraded-but-safe* - conditions (empty runtime, malformed existing config), which warn + skip.""" + """A resolved asset could not be applied. Raised so the task fails-closed + rather than running with an audit record claiming an asset that was never + actually loaded (#246 Option C). This covers every condition that would leave + a pinned asset unloaded — missing repo dir, empty/invalid runtime, structurally + invalid connection config, or a write error (see :func:`apply_mcp_assets`). The + only warn-and-continue case is a pre-existing malformed ``.mcp.json`` on disk + (:func:`_read_existing_mcp_config`), which is replaced, not fatal.""" + + +def _protect_mcp_json_from_commit(repo_dir: str, mcp_path: str) -> None: + """Mark ``.mcp.json`` skip-worktree so the safety-net commit can't push it. + + The resolved runtime we just wrote may carry secret-bearing fields (bearer + ``headers``, a ``url`` with a token query string, ``args`` like + ``--api-key=…``). ``.mcp.json`` lives in the live git clone, and the + post-hook safety net (``git add -u`` → commit → ``git push``) will exfiltrate + it into the PR whenever the target repo *tracks* ``.mcp.json`` (#246 review + B4). Setting ``skip-worktree`` blocks both ``git add -u`` and an explicit + ``git add .mcp.json``; for an untracked file we first ``--intent-to-add`` so + the flag has an index entry to attach to. Mechanical enforcement, mirroring + the ADR-016 Linear-strip posture — a prompt is not a security boundary. + + Best-effort: git plumbing failures here are logged, not fatal. The SDK still + reads the on-disk file (skip-worktree only affects the index, not the + working tree), so loading is unaffected either way. + """ + + def _git(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["git", "-C", repo_dir, *args], + capture_output=True, + text=True, + timeout=_GIT_GUARD_TIMEOUT_S, + check=False, + ) + + try: + tracked = _git("ls-files", "--error-unmatch", ".mcp.json").returncode == 0 + if not tracked: + # Give the untracked file an index entry so skip-worktree can attach. + _git("add", "--intent-to-add", ".mcp.json") + result = _git("update-index", "--skip-worktree", ".mcp.json") + if result.returncode != 0: + log( + "WARN", + f"Registry: could not skip-worktree {mcp_path} " + f"(exit {result.returncode}): {result.stderr.strip()}", + ) + except (OSError, subprocess.SubprocessError) as e: + log("WARN", f"Registry: skip-worktree guard for {mcp_path} failed: {type(e).__name__}: {e}") def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> list[str]: @@ -155,9 +204,6 @@ def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> li servers[key] = _to_mcp_config(runtime, key) written.append(key) - if not written: - return [] - config["mcpServers"] = servers try: with open(mcp_path, "w", encoding="utf-8") as f: @@ -166,6 +212,10 @@ def apply_mcp_assets(repo_dir: str, resolved_assets: list[dict[str, Any]]) -> li except OSError as e: raise RegistryAssetLoadError(f"failed to write {mcp_path}: {e}") from e + # The runtime we just wrote may carry secrets; keep the safety-net commit + # from pushing it to the PR (#246 review B4). + _protect_mcp_json_from_commit(repo_dir, mcp_path) + log("TASK", f"Registry: merged {len(written)} MCP server(s) into {mcp_path}") return written diff --git a/agent/tests/test_registry_loader.py b/agent/tests/test_registry_loader.py index ed06ad831..21717ecbb 100644 --- a/agent/tests/test_registry_loader.py +++ b/agent/tests/test_registry_loader.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import subprocess import pytest @@ -284,3 +285,64 @@ def test_registry_linear_server_is_stripped_after_merge(self, tmp_path): servers = _read_mcp(tmp_path)["mcpServers"] assert "evil__linear" not in servers # Linear scrubbed assert "acme__pdf" in servers # benign server survives + + +class TestMcpJsonNotCommittable: + """A written .mcp.json carries the resolved runtime, which may hold secrets + (bearer headers, url tokens, --api-key args). It lives in the live git clone, + so the post-hook `git add -u` → commit → push must NOT be able to exfiltrate + it to the PR. apply_mcp_assets marks it skip-worktree to block that (#246 B4).""" + + @staticmethod + def _git(repo, *args) -> subprocess.CompletedProcess: + return subprocess.run( + ["git", "-C", str(repo), *args], + capture_output=True, + text=True, + check=False, + ) + + def _init_repo(self, tmp_path): + self._git(tmp_path, "init", "-q") + self._git(tmp_path, "config", "user.email", "t@t") + self._git(tmp_path, "config", "user.name", "t") + (tmp_path / "README.md").write_text("x") + self._git(tmp_path, "add", "README.md") + self._git(tmp_path, "commit", "-qm", "init") + + def _secret_asset(self): + return _mcp_asset( + "acme", + "secretful", + "1.0.0", + { + "transport": "http", + "url": "https://mcp.example/sse?token=SUPERSECRET", + "headers": {"Authorization": "Bearer sk-live-abc123"}, + }, + ) + + def test_untracked_mcp_json_cannot_be_staged(self, tmp_path): + self._init_repo(tmp_path) + apply_mcp_assets(str(tmp_path), [self._secret_asset()]) + # The file exists on disk (the SDK still reads it)... + assert (tmp_path / ".mcp.json").exists() + # ...but the safety-net `git add -u` (and even an explicit add) cannot stage it. + self._git(tmp_path, "add", "-u") + self._git(tmp_path, "add", ".mcp.json") + staged = self._git(tmp_path, "diff", "--cached") + assert "SUPERSECRET" not in staged.stdout + assert "sk-live-abc123" not in staged.stdout + + def test_tracked_mcp_json_change_cannot_be_staged(self, tmp_path): + # The dangerous case Scott reproduced: the repo already tracks .mcp.json. + self._init_repo(tmp_path) + (tmp_path / ".mcp.json").write_text('{"mcpServers":{}}\n') + self._git(tmp_path, "add", ".mcp.json") + self._git(tmp_path, "commit", "-qm", "track mcp") + apply_mcp_assets(str(tmp_path), [self._secret_asset()]) + # git add -u stages tracked-but-modified files — must skip .mcp.json now. + self._git(tmp_path, "add", "-u") + staged = self._git(tmp_path, "diff", "--cached") + assert staged.stdout.strip() == "" + assert "SUPERSECRET" not in staged.stdout diff --git a/cdk/src/constructs/task-orchestrator.ts b/cdk/src/constructs/task-orchestrator.ts index e21307e4e..47fb673ed 100644 --- a/cdk/src/constructs/task-orchestrator.ts +++ b/cdk/src/constructs/task-orchestrator.ts @@ -474,8 +474,10 @@ export class TaskOrchestrator extends Construct { })); // 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. + // Blueprint's registry:// asset refs at task start. Scoped to THIS registry + // (the id is in scope here); only the record suffix is a wildcard, because + // record ids are server-assigned and unknown at synth — mirrors the scoping + // in registry-api.ts (#246 review nit). if (props.agentRegistryId) { this.fn.addToRolePolicy(new iam.PolicyStatement({ actions: [ @@ -486,13 +488,13 @@ export class TaskOrchestrator extends Construct { Stack.of(this).formatArn({ service: 'bedrock-agentcore', resource: 'registry', - resourceName: '*', + resourceName: props.agentRegistryId, arnFormat: ArnFormat.SLASH_RESOURCE_NAME, }), Stack.of(this).formatArn({ service: 'bedrock-agentcore', resource: 'registry', - resourceName: '*/record/*', + resourceName: `${props.agentRegistryId}/record/*`, arnFormat: ArnFormat.SLASH_RESOURCE_NAME, }), ], @@ -662,7 +664,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; S3 object/* wildcard from CDK grantPut on the dedicated MicroVM payload bucket; MicroVM lifecycle actions (RunMicrovm/GetMicrovm/TerminateMicrovm) are scoped to the single platform MicroVM image ARN plus a :* version-suffix sibling (every one of them authorizes against the image resource, not the per-session instance; no account-wide wildcard is used); lambda:PassNetworkConnector requires Resource:* because the action supports no resource-level permissions and the AWS-managed connectors live outside this account; iam:PassRole is scoped to the MicroVM execution role and conditioned on lambda.amazonaws.com; AgentCore registry/* + registry/*/record/* wildcards because record ids are server-assigned (#246)', + 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; S3 object/* wildcard from CDK grantPut on the dedicated MicroVM payload bucket; MicroVM lifecycle actions (RunMicrovm/GetMicrovm/TerminateMicrovm) are scoped to the single platform MicroVM image ARN plus a :* version-suffix sibling (every one of them authorizes against the image resource, not the per-session instance; no account-wide wildcard is used); lambda:PassNetworkConnector requires Resource:* because the action supports no resource-level permissions and the AWS-managed connectors live outside this account; iam:PassRole is scoped to the MicroVM execution role and conditioned on lambda.amazonaws.com; AgentCore registry read scoped to the wired registry ARN, with a record/* suffix wildcard because record ids are server-assigned and unknown at synth (#246)', }, ], true); } diff --git a/cdk/test/constructs/task-orchestrator.test.ts b/cdk/test/constructs/task-orchestrator.test.ts index 418617928..829171f69 100644 --- a/cdk/test/constructs/task-orchestrator.test.ts +++ b/cdk/test/constructs/task-orchestrator.test.ts @@ -43,6 +43,7 @@ interface StackOverrides { taskRoleArn: string; executionRoleArn: string; }; + agentRegistryId?: string; } function createStack(overrides?: StackOverrides): { stack: Stack; template: Template } { @@ -453,6 +454,29 @@ describe('TaskOrchestrator construct', () => { ); }); + test('registry read grant is scoped to the wired registry id — no bare "*" (#246 review)', () => { + const { template } = createStack({ agentRegistryId: 'reg-xyz789' }); + const policies = template.findResources('AWS::IAM::Policy'); + const serialized = JSON.stringify(policies); + // The wired registry id must appear in the resource ARNs... + expect(serialized).toContain('reg-xyz789'); + // ...and no bedrock-agentcore registry statement may grant a bare "*" or a + // registry/* wildcard (the finding: it should scope to registry/{id}). + for (const policy of Object.values(policies)) { + const statements = (policy as { + Properties: { PolicyDocument: { Statement: Array<{ Action: unknown; Resource: unknown }> } }; + }).Properties.PolicyDocument.Statement; + for (const stmt of statements) { + const actions = JSON.stringify(stmt.Action); + if (actions.includes('bedrock-agentcore:GetRegistryRecord')) { + const resources = JSON.stringify(stmt.Resource); + expect(resources).not.toContain('registry/*'); + expect(stmt.Resource).not.toBe('*'); + } + } + } + }); + describe('ECS compute strategy', () => { test('includes ECS env vars when ECS props are provided', () => { ecsTemplate.hasResourceProperties('AWS::Lambda::Function', {