Merge branch 'main' into feat/browser-screenshot-integration

Author: krokoko

agent/src/memory.py

@@ -7,18 +7,27 @@
 ERROR level to surface bugs quickly.
 """
 
+import hashlib
 import os
 import re
 import time
 
+from sanitization import sanitize_external_content
+
 _client = None
 
 # Validates "owner/repo" format — must match the TypeScript-side isValidRepo pattern.
 _REPO_PATTERN = re.compile(r"^[a-zA-Z0-9._-]+/[a-zA-Z0-9._-]+$")
 
-# Current event schema version — used to distinguish records written under
-# different namespace schemes (v1 = repos/ prefix, v2 = namespace templates).
-_SCHEMA_VERSION = "2"
+# Current event schema version:
+#   v1 = repos/ prefix
+#   v2 = namespace templates (/{actorId}/...)
+#   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():
@@ -50,7 +59,8 @@ def _log_error(func_name: str, err: Exception, memory_id: str, task_id: str) ->
     level = "ERROR" if is_programming_error else "WARN"
     label = "unexpected error" if is_programming_error else "infra failure"
     print(
-        f"[memory] [{level}] {func_name} {label}: {type(err).__name__}",
+        f"[memory] [{level}] {func_name} {label}: {type(err).__name__}: {err}"
+        f" (memory_id={memory_id}, task_id={task_id})",
         flush=True,
     )
 
@@ -75,6 +85,9 @@ def write_task_episode(
     namespace templates (/{actorId}/episodes/{sessionId}/) place records
     into the correct per-repo, per-task namespace.
 
+    Metadata includes source_type='agent_episode' for provenance tracking
+    and content_sha256 for integrity auditing on read (schema v3).
+
     Returns True on success, False on failure (fail-open).
     """
     try:
@@ -94,10 +107,16 @@ def write_task_episode(
             parts.append(f"Agent notes: {self_feedback}")
 
         episode_text = " ".join(parts)
+        # 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},
             "type": {"stringValue": "task_episode"},
+            "source_type": {"stringValue": "agent_episode"},
+            "content_sha256": {"stringValue": content_hash},
             "schema_version": {"stringValue": _SCHEMA_VERSION},
         }
         if pr_url:
@@ -142,12 +161,24 @@ def write_repo_learnings(
     namespace templates (/{actorId}/knowledge/) place records into
     the correct per-repo namespace.
 
+    Metadata includes source_type='agent_learning' for provenance tracking
+    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).
     """
     try:
         _validate_repo(repo)
         client = _get_client()
 
+        learnings_text = f"Repository learnings: {learnings}"
+        # 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,
             actorId=repo,
@@ -156,14 +187,16 @@ def write_repo_learnings(
             payload=[
                 {
                     "conversational": {
-                        "content": {"text": f"Repository learnings: {learnings}"},
+                        "content": {"text": learnings_text},
                         "role": "OTHER",
                     }
                 }
             ],
             metadata={
                 "task_id": {"stringValue": task_id},
                 "type": {"stringValue": "repo_learnings"},
+                "source_type": {"stringValue": "agent_learning"},
+                "content_sha256": {"stringValue": content_hash},
                 "schema_version": {"stringValue": _SCHEMA_VERSION},
             },
         )

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

@@ -8,6 +8,7 @@
 
 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
 
@@ -49,11 +50,11 @@ def build_system_prompt(
         if mc.repo_knowledge:
             mc_parts.append("**Repository knowledge:**")
             for item in mc.repo_knowledge:
-                mc_parts.append(f"- {item}")
+                mc_parts.append(f"- {sanitize_memory_content(item)}")
         if mc.past_episodes:
             mc_parts.append("\n**Past task episodes:**")
             for item in mc.past_episodes:
-                mc_parts.append(f"- {item}")
+                mc_parts.append(f"- {sanitize_memory_content(item)}")
         if mc_parts:
             memory_context_text = "\n".join(mc_parts)
     system_prompt = system_prompt.replace("{memory_context}", memory_context_text)

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,8 +1,18 @@
 """Unit tests for pure functions in memory.py."""
 
+import hashlib
+from unittest.mock import MagicMock, patch
+
 import pytest
 
-from memory import _validate_repo
+from memory import (
+    _SCHEMA_VERSION,
+    MEMORY_SOURCE_TYPES,
+    _validate_repo,
+    write_repo_learnings,
+    write_task_episode,
+)
+from sanitization import sanitize_external_content
 
 
 class TestValidateRepo:
@@ -34,3 +44,83 @@ def test_invalid_spaces(self):
     def test_invalid_empty(self):
         with pytest.raises(ValueError, match="does not match"):
             _validate_repo("")
+
+
+class TestSchemaVersion:
+    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):
+        mock_client = MagicMock()
+        mock_get_client.return_value = mock_client
+
+        write_task_episode("mem-1", "owner/repo", "task-1", "COMPLETED")
+
+        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_content_sha256_matches_sanitized_content(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_get_client.return_value = mock_client
+
+        write_task_episode("mem-1", "owner/repo", "task-1", "COMPLETED")
+
+        call_kwargs = mock_client.create_event.call_args[1]
+        metadata = call_kwargs["metadata"]
+        assert "content_sha256" in metadata
+        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:
+    @patch("memory._get_client")
+    def test_includes_source_type_in_metadata(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_get_client.return_value = mock_client
+
+        write_repo_learnings("mem-1", "owner/repo", "task-1", "Use Jest for tests")
+
+        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_content_sha256_matches_sanitized_content(self, mock_get_client):
+        mock_client = MagicMock()
+        mock_get_client.return_value = mock_client
+
+        write_repo_learnings("mem-1", "owner/repo", "task-1", "Use Jest for tests")
+
+        call_kwargs = mock_client.create_event.call_args[1]
+        metadata = call_kwargs["metadata"]
+        assert "content_sha256" in metadata
+        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):