Merge branch 'main' into feat/ec2-fleet-strategy

Author: MichaelWalker-git

agent/src/memory.py

@@ -12,6 +12,8 @@
 import re
 import time
 
+from sanitization import sanitize_external_content
+
 _client = None
 
 # Validates "owner/repo" format — must match the TypeScript-side isValidRepo pattern.
@@ -23,6 +25,10 @@
 #   v3 = adds source_type provenance + content_sha256 integrity hash
 _SCHEMA_VERSION = "3"
 
+# Valid source_type values for provenance tracking (schema v3).
+# Must stay in sync with MemorySourceType in cdk/src/handlers/shared/memory.ts.
+MEMORY_SOURCE_TYPES = frozenset({"agent_episode", "agent_learning", "orchestrator_fallback"})
+
 
 def _get_client():
     """Lazy-init and cache the AgentCore client for memory operations."""
@@ -80,7 +86,7 @@ def write_task_episode(
     into the correct per-repo, per-task namespace.
 
     Metadata includes source_type='agent_episode' for provenance tracking
-    and content_sha256 for integrity verification on read (schema v3).
+    and content_sha256 for integrity auditing on read (schema v3).
 
     Returns True on success, False on failure (fail-open).
     """
@@ -101,7 +107,10 @@ def write_task_episode(
             parts.append(f"Agent notes: {self_feedback}")
 
         episode_text = " ".join(parts)
-        content_hash = hashlib.sha256(episode_text.encode("utf-8")).hexdigest()
+        # Hash the sanitized form; store the original. The read path re-sanitizes
+        # and checks against this hash: sanitize(original) at write == sanitize(stored) at read.
+        sanitized_text = sanitize_external_content(episode_text)
+        content_hash = hashlib.sha256(sanitized_text.encode("utf-8")).hexdigest()
 
         metadata = {
             "task_id": {"stringValue": task_id},
@@ -153,9 +162,10 @@ def write_repo_learnings(
     the correct per-repo namespace.
 
     Metadata includes source_type='agent_learning' for provenance tracking
-    and content_sha256 for integrity verification on read (schema v3).
-    Note: hash verification only happens on the TS orchestrator read path
-    (loadMemoryContext in memory.ts), not on the Python side.
+    and content_sha256 for integrity auditing on read (schema v3).
+    Note: hash auditing only happens on the TS orchestrator read path
+    (loadMemoryContext in memory.ts) where mismatches are logged but
+    records are kept — the Python side does not independently check hashes.
 
     Returns True on success, False on failure (fail-open).
     """
@@ -164,7 +174,10 @@ def write_repo_learnings(
         client = _get_client()
 
         learnings_text = f"Repository learnings: {learnings}"
-        content_hash = hashlib.sha256(learnings_text.encode("utf-8")).hexdigest()
+        # Hash the sanitized form; store the original. The read path re-sanitizes
+        # and checks against this hash: sanitize(original) at write == sanitize(stored) at read.
+        sanitized_text = sanitize_external_content(learnings_text)
+        content_hash = hashlib.sha256(sanitized_text.encode("utf-8")).hexdigest()
 
         client.create_event(
             memoryId=memory_id,

agent/src/models.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from enum import StrEnum
-from typing import Self
+from typing import Literal, Self
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
@@ -52,6 +52,11 @@ class MemoryContext(BaseModel):
     past_episodes: list[str] = Field(default_factory=list)
 
 
+# Trust classification for content sources — mirrors ContentTrustLevel in context-hydration.ts.
+# 'trusted': user-supplied input, 'untrusted-external': GitHub-sourced content,
+# 'memory': memory records.
+ContentTrustLevel = Literal["trusted", "untrusted-external", "memory"]
+
 # Bump when this agent supports a new orchestrator HydratedContext shape
 # (see cdk/src/handlers/shared/context-hydration.ts).
 SUPPORTED_HYDRATED_CONTEXT_VERSION = 1
@@ -73,6 +78,7 @@ class HydratedContext(BaseModel):
     guardrail_blocked: str | None = None
     resolved_branch_name: str | None = None
     resolved_base_branch: str | None = None
+    content_trust: dict[str, ContentTrustLevel] | None = None
 
     @model_validator(mode="after")
     def version_supported(self) -> Self:

agent/src/prompt_builder.py

@@ -4,53 +4,14 @@
 
 import glob
 import os
-import re
 from typing import TYPE_CHECKING
 
 from config import AGENT_WORKSPACE
 from prompts import get_system_prompt
+from sanitization import sanitize_external_content as sanitize_memory_content
 from shell import log
 from system_prompt import SYSTEM_PROMPT
 
-# ---------------------------------------------------------------------------
-# Content sanitization for memory records
-# ---------------------------------------------------------------------------
-
-_DANGEROUS_TAGS = re.compile(
-    r"(<(script|style|iframe|object|embed|form|input)[^>]*>[\s\S]*?</\2>"
-    r"|<(script|style|iframe|object|embed|form|input)[^>]*\/?>)",
-    re.IGNORECASE,
-)
-_HTML_TAGS = re.compile(r"</?[a-z][^>]*>", re.IGNORECASE)
-_INSTRUCTION_PREFIXES = re.compile(
-    r"^(SYSTEM|ASSISTANT|Human|Assistant)\s*:", re.MULTILINE | re.IGNORECASE
-)
-_INJECTION_PHRASES = re.compile(
-    r"(?:ignore previous instructions|disregard (?:above|previous|all)|new instructions\s*:)",
-    re.IGNORECASE,
-)
-_CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]")
-_BIDI_CHARS = re.compile(r"[\u200e\u200f\u202a-\u202e\u2066-\u2069]")
-_MISPLACED_BOM = re.compile(r"(?!^)\ufeff")
-
-
-def sanitize_memory_content(text: str | None) -> str:
-    """Sanitize memory content before injecting into the agent's system prompt.
-
-    Mirrors the TypeScript sanitizeExternalContent() in sanitization.ts.
-    """
-    if not text:
-        return text or ""
-    s = _DANGEROUS_TAGS.sub("", text)
-    s = _HTML_TAGS.sub("", s)
-    s = _INSTRUCTION_PREFIXES.sub(r"[SANITIZED_PREFIX] \1:", s)
-    s = _INJECTION_PHRASES.sub("[SANITIZED_INSTRUCTION]", s)
-    s = _CONTROL_CHARS.sub("", s)
-    s = _BIDI_CHARS.sub("", s)
-    s = _MISPLACED_BOM.sub("", s)
-    return s
-
-
 if TYPE_CHECKING:
     from models import HydratedContext, RepoSetup, TaskConfig
 

agent/src/sanitization.py

@@ -0,0 +1,60 @@
+"""Content sanitization for external/untrusted inputs.
+
+Mirrors the TypeScript sanitizeExternalContent() in
+cdk/src/handlers/shared/sanitization.ts. Both implementations
+must produce identical output for the same input — cross-language
+parity is verified by shared test fixtures.
+
+Applied to: memory records (before hashing on write, before injection
+on read), GitHub issue/PR content (TS side only — Python agent receives
+already-sanitized content from the orchestrator's hydrated context).
+"""
+
+import re
+
+_DANGEROUS_TAGS = re.compile(
+    r"(<(script|style|iframe|object|embed|form|input)[^>]*>[\s\S]*?</\2>"
+    r"|<(script|style|iframe|object|embed|form|input)[^>]*\/?>)",
+    re.IGNORECASE,
+)
+_HTML_TAGS = re.compile(r"</?[a-z][^>]*>", re.IGNORECASE)
+_INSTRUCTION_PREFIXES = re.compile(r"^(SYSTEM|ASSISTANT|Human)\s*:", re.MULTILINE | re.IGNORECASE)
+_INJECTION_PHRASES = re.compile(
+    r"(?:ignore previous instructions|disregard (?:above|previous|all)|new instructions\s*:)",
+    re.IGNORECASE,
+)
+_CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]")
+_BIDI_CHARS = re.compile(r"[\u200e\u200f\u202a-\u202e\u2066-\u2069]")
+_MISPLACED_BOM = re.compile(r"(?!^)\ufeff")
+
+
+def _strip_until_stable(s: str, pattern: re.Pattern[str]) -> str:
+    """Apply *pattern* repeatedly until the string stops changing.
+
+    A single pass can be bypassed by nesting fragments
+    (e.g. "<scrip<script></script>t>" reassembles after inner tag removal).
+    """
+    while True:
+        prev = s
+        s = pattern.sub("", s)
+        if s == prev:
+            return s
+
+
+def sanitize_external_content(text: str | None) -> str:
+    """Sanitize external content before it enters the agent's context.
+
+    Neutralizes rather than blocks — suspicious patterns are replaced with
+    bracketed markers so content is still visible to the LLM (for legitimate
+    discussion of prompts/instructions) but structurally defanged.
+    """
+    if not text:
+        return text or ""
+    s = _strip_until_stable(text, _DANGEROUS_TAGS)
+    s = _strip_until_stable(s, _HTML_TAGS)
+    s = _INSTRUCTION_PREFIXES.sub(r"[SANITIZED_PREFIX] \1:", s)
+    s = _INJECTION_PHRASES.sub("[SANITIZED_INSTRUCTION]", s)
+    s = _CONTROL_CHARS.sub("", s)
+    s = _BIDI_CHARS.sub("", s)
+    s = _MISPLACED_BOM.sub("", s)
+    return s

agent/tests/test_memory.py

@@ -1,10 +1,18 @@
 """Unit tests for pure functions in memory.py."""
 
+import hashlib
 from unittest.mock import MagicMock, patch
 
 import pytest
 
-from memory import _SCHEMA_VERSION, _validate_repo, write_repo_learnings, write_task_episode
+from memory import (
+    _SCHEMA_VERSION,
+    MEMORY_SOURCE_TYPES,
+    _validate_repo,
+    write_repo_learnings,
+    write_task_episode,
+)
+from sanitization import sanitize_external_content
 
 
 class TestValidateRepo:
@@ -43,6 +51,14 @@ def test_schema_version_is_3(self):
         assert _SCHEMA_VERSION == "3"
 
 
+class TestMemorySourceTypes:
+    def test_contains_expected_values(self):
+        assert {"agent_episode", "agent_learning", "orchestrator_fallback"} == MEMORY_SOURCE_TYPES
+
+    def test_is_frozen(self):
+        assert isinstance(MEMORY_SOURCE_TYPES, frozenset)
+
+
 class TestWriteTaskEpisode:
     @patch("memory._get_client")
     def test_includes_source_type_in_metadata(self, mock_get_client):
@@ -54,10 +70,11 @@ def test_includes_source_type_in_metadata(self, mock_get_client):
         call_kwargs = mock_client.create_event.call_args[1]
         metadata = call_kwargs["metadata"]
         assert metadata["source_type"] == {"stringValue": "agent_episode"}
+        assert metadata["source_type"]["stringValue"] in MEMORY_SOURCE_TYPES
         assert metadata["schema_version"] == {"stringValue": "3"}
 
     @patch("memory._get_client")
-    def test_includes_content_sha256_in_metadata(self, mock_get_client):
+    def test_content_sha256_matches_sanitized_content(self, mock_get_client):
         mock_client = MagicMock()
         mock_get_client.return_value = mock_client
 
@@ -66,8 +83,14 @@ def test_includes_content_sha256_in_metadata(self, mock_get_client):
         call_kwargs = mock_client.create_event.call_args[1]
         metadata = call_kwargs["metadata"]
         assert "content_sha256" in metadata
-        # SHA-256 hex is 64 chars
-        assert len(metadata["content_sha256"]["stringValue"]) == 64
+        hash_value = metadata["content_sha256"]["stringValue"]
+        assert len(hash_value) == 64
+
+        # Verify hash matches the sanitized content that was actually stored
+        content = call_kwargs["payload"][0]["conversational"]["content"]["text"]
+        sanitized = sanitize_external_content(content)
+        expected = hashlib.sha256(sanitized.encode("utf-8")).hexdigest()
+        assert hash_value == expected
 
 
 class TestWriteRepoLearnings:
@@ -81,10 +104,11 @@ def test_includes_source_type_in_metadata(self, mock_get_client):
         call_kwargs = mock_client.create_event.call_args[1]
         metadata = call_kwargs["metadata"]
         assert metadata["source_type"] == {"stringValue": "agent_learning"}
+        assert metadata["source_type"]["stringValue"] in MEMORY_SOURCE_TYPES
         assert metadata["schema_version"] == {"stringValue": "3"}
 
     @patch("memory._get_client")
-    def test_includes_content_sha256_in_metadata(self, mock_get_client):
+    def test_content_sha256_matches_sanitized_content(self, mock_get_client):
         mock_client = MagicMock()
         mock_get_client.return_value = mock_client
 
@@ -93,4 +117,10 @@ def test_includes_content_sha256_in_metadata(self, mock_get_client):
         call_kwargs = mock_client.create_event.call_args[1]
         metadata = call_kwargs["metadata"]
         assert "content_sha256" in metadata
-        assert len(metadata["content_sha256"]["stringValue"]) == 64
+        hash_value = metadata["content_sha256"]["stringValue"]
+        assert len(hash_value) == 64
+
+        content = call_kwargs["payload"][0]["conversational"]["content"]["text"]
+        sanitized = sanitize_external_content(content)
+        expected = hashlib.sha256(sanitized.encode("utf-8")).hexdigest()
+        assert hash_value == expected

agent/tests/test_models.py

@@ -201,6 +201,47 @@ def test_extra_top_level_forbidden(self):
                 }
             )
 
+    def test_content_trust_none_by_default(self):
+        hc = HydratedContext(user_prompt="Fix bug")
+        assert hc.content_trust is None
+
+    def test_content_trust_accepted(self):
+        hc = HydratedContext(
+            user_prompt="Fix bug",
+            content_trust={"issue": "untrusted-external", "task_description": "trusted"},
+        )
+        assert hc.content_trust == {"issue": "untrusted-external", "task_description": "trusted"}
+
+    def test_content_trust_with_memory(self):
+        hc = HydratedContext(
+            user_prompt="Fix bug",
+            content_trust={"memory": "memory", "task_description": "trusted"},
+        )
+        assert hc.content_trust is not None
+        assert hc.content_trust["memory"] == "memory"
+
+    def test_content_trust_round_trip(self):
+        data = {
+            "version": 1,
+            "user_prompt": "Do the thing",
+            "sources": ["issue", "memory"],
+            "content_trust": {
+                "issue": "untrusted-external",
+                "memory": "memory",
+            },
+        }
+        hc = HydratedContext.model_validate(data)
+        assert hc.content_trust == {"issue": "untrusted-external", "memory": "memory"}
+
+    def test_content_trust_invalid_value_rejected(self):
+        with pytest.raises(ValidationError):
+            HydratedContext.model_validate(
+                {
+                    "user_prompt": "Fix bug",
+                    "content_trust": {"issue": "invalid-trust-level"},
+                }
+            )
+
 
 class TestTaskConfig:
     def test_required_fields(self):