Redact a job body wherever it sits, and stop three crashes

.changeset/redaction-robustness.md

@@ -0,0 +1,5 @@
+---
+"apollo": patch
+---
+
+Replace a job body wherever it sits in the workflow, not only at the top level, and stop three crashes on a shape the model did not expect

services/global_chat/tests/unit/test_redaction_robustness.py

@@ -0,0 +1,109 @@
+"""Redaction has to survive an odd document rather than throw or mangle it.
+
+Bodies are deferred to save tokens, not hidden: the model can read any of them
+with inspect_job_code. So a document we cannot parse comes back as it came,
+where the model can diagnose it. A YAML anchor can point at its own container,
+which PyYAML builds as a real cycle.
+"""
+
+import pytest
+import yaml
+from workflow_chat.workflow_chat import AnthropicClient
+from yaml_utils import (
+    _remove_ids,
+    redact_job_bodies,
+    workflow_has_job_code,
+)
+
+SECRET = "callSecretApi()"
+
+
+def test_a_normal_workflow_is_still_redacted() -> None:
+    out = redact_job_bodies(f"jobs:\n  a:\n    id: 123\n    body: {SECRET}\n")
+
+    assert SECRET not in out
+    assert "123" not in out
+    assert "jobs" in out
+
+
+@pytest.mark.parametrize(
+    "document",
+    [
+        f'jobs: {{a: {{body: "{SECRET}"}}\n  broken',
+        f"- {SECRET}\n",
+        f"jobs: {SECRET}\n",
+    ],
+    ids=["unparseable", "a-list", "jobs-not-a-mapping"],
+)
+def test_a_document_it_cannot_redact_comes_back_as_it_came(document: str) -> None:
+    """Withholding it would only cost the model the structure. It can already
+    read any body it wants, so there is nothing here to keep from it."""
+    assert redact_job_bodies(document) == document
+
+
+def test_the_id_walk_terminates_on_a_self_referential_anchor() -> None:
+    data = yaml.safe_load("jobs:\n  a: &x\n    body: code()\n    loop: *x\n")
+    assert data["jobs"]["a"]["loop"] is data["jobs"]["a"]
+
+    _remove_ids(data)
+
+    assert "id" not in data["jobs"]["a"]
+
+
+def test_redaction_terminates_on_a_self_referential_anchor() -> None:
+    out = redact_job_bodies(f"jobs:\n  a: &x\n    body: {SECRET}\n    loop: *x\n")
+
+    assert SECRET not in out
+
+
+@pytest.mark.parametrize(
+    "yaml_data",
+    [
+        {"jobs": [{"body": "code()"}]},
+        {"jobs": {"a": {"body": 42}}},
+        {"jobs": {"a": None}},
+        {"jobs": "not a mapping"},
+        {"triggers": {"t": None}},
+        {"edges": {"e": "not a mapping"}},
+    ],
+    ids=["jobs-a-list", "numeric-body", "null-job", "jobs-a-string",
+         "null-trigger", "edge-not-a-mapping"],
+)
+def test_preserving_components_tolerates_a_shape_it_did_not_expect(yaml_data: dict) -> None:
+    preserved, _ = AnthropicClient.extract_and_preserve_components(yaml_data)
+
+    assert isinstance(preserved, dict)
+
+@pytest.mark.parametrize(
+    "document",
+    [
+        f'jobs:\n  a:\n    body: "ok"\n    steps:\n      inner:\n        body: "{SECRET}"\n',
+        f'jobs:\n  a:\n    - body: "{SECRET}"\n',
+        f'shared: &s\n  body: "{SECRET}"\njobs:\n  a:\n    <<: *s\n',
+        f"workflows:\n  wf:\n    jobs:\n      a:\n        body: {SECRET}\n",
+    ],
+    ids=["nested-deeper", "job-is-a-list", "merge-key", "project-export"],
+)
+def test_a_body_is_redacted_wherever_it_sits(document: str) -> None:
+    out = redact_job_bodies(document)
+
+    assert SECRET not in out
+
+
+def test_a_workflow_with_no_bodies_is_kept() -> None:
+    out = redact_job_bodies("triggers:\n  t:\n    type: cron\n")
+
+    assert "cron" in out
+
+
+def test_the_read_only_id_strip_also_survives_a_cycle() -> None:
+    document = "jobs:\n  a: &x\n    id: SECRET-ID\n    body: code()\n    loop: *x\n"
+
+    out = AnthropicClient.remove_ids_from_yaml(AnthropicClient, document)
+
+    assert "SECRET-ID" not in out
+
+
+def test_a_scalar_document_is_not_treated_as_a_workflow() -> None:
+    assert redact_job_bodies("jobs") == "jobs"
+    assert workflow_has_job_code("jobs") is False

services/workflow_chat/workflow_chat.py

@@ -63,6 +63,7 @@
 from langfuse import observe, propagate_attributes, get_client as get_langfuse_client
 from langfuse_util import should_track, build_tags, build_generation_diff, mask_secrets
 from util import ApolloError, create_logger, add_page_prefix, APOLLO_VERSION
+from yaml_utils import _remove_ids
 from .gen_project_prompt import build_prompt
 from workflow_chat.available_adaptors import get_available_adaptors
 from streaming_util import (
@@ -373,17 +374,8 @@ def remove_ids_from_yaml(self, yaml_str):
             return yaml_str
         try:
             yaml_data = yaml.safe_load(yaml_str)
-
-            def remove_ids(obj):
-                if isinstance(obj, dict):
-                    obj.pop("id", None)
-                    for v in obj.values():
-                        remove_ids(v)
-                elif isinstance(obj, list):
-                    for item in obj:
-                        remove_ids(item)
-
-            remove_ids(yaml_data)
+            # Shared with redact_job_bodies so the cycle guard lives in one place.
+            _remove_ids(yaml_data)
             return yaml.dump(yaml_data, sort_keys=False, default_flow_style=False)
         except Exception as e:
             logger.warning(f"Could not remove IDs from YAML: {e}")
@@ -577,9 +569,11 @@ def extract_and_preserve_components(yaml_data):
 
         preserved_values = {}
 
-        if "jobs" in yaml_data:
+        if isinstance(yaml_data.get("jobs"), dict):
             for job_key, job_data in yaml_data["jobs"].items():
-                if "body" in job_data:
+                if not isinstance(job_data, dict):
+                    continue
+                if isinstance(job_data.get("body"), str):
                     body_content = job_data["body"].strip()
                     if body_content and body_content != "// Add operations here":
                         placeholder = f"__CODE_BLOCK_{job_key}__"
@@ -591,18 +585,18 @@ def extract_and_preserve_components(yaml_data):
                     preserved_values[placeholder] = job_data["id"]
                     job_data["id"] = placeholder
 
-        if "triggers" in yaml_data:
+        if isinstance(yaml_data.get("triggers"), dict):
             for trigger_key, trigger_data in yaml_data["triggers"].items():
-                if "id" in trigger_data:
+                if isinstance(trigger_data, dict) and "id" in trigger_data:
                     # Flat, not keyed on the trigger name: the model renames
                     # that key when it swaps webhook for cron.
                     preserved_values["trigger_id"] = trigger_data["id"]
                     # Remove the id key from what we send to the model
                     del trigger_data["id"]
 
-        if "edges" in yaml_data:
+        if isinstance(yaml_data.get("edges"), dict):
             for edge_key, edge_data in yaml_data["edges"].items():
-                if "id" in edge_data:
+                if isinstance(edge_data, dict) and "id" in edge_data:
                     placeholder = f"__ID_EDGE_{edge_key}__"
                     preserved_values[placeholder] = edge_data["id"]
                     edge_data["id"] = placeholder

services/yaml_utils.py

@@ -69,7 +69,7 @@ def find_job_in_yaml(yaml_str: str, step_name: str) -> tuple[str | None, dict |
     except Exception:
         return None, None
 
-    if not yaml_data or "jobs" not in yaml_data:
+    if not isinstance(yaml_data, dict) or "jobs" not in yaml_data:
         return None, None
 
     jobs = yaml_data["jobs"]
@@ -105,7 +105,7 @@ def workflow_has_job_code(yaml_str: str | None) -> bool:
         yaml_data = yaml.safe_load(yaml_str)
     except Exception:
         return False
-    if not yaml_data or "jobs" not in yaml_data:
+    if not isinstance(yaml_data, dict) or not isinstance(yaml_data.get("jobs"), dict):
         return False
     for job_data in yaml_data["jobs"].values():
         body = (job_data or {}).get("body")
@@ -114,36 +114,76 @@ def workflow_has_job_code(yaml_str: str | None) -> bool:
     return False
 
 
+#: What a job body is replaced with in the structural view.
+REDACTED_BODY = "# [use inspect_job_code to view]"
+
+
+def _redact_bodies(obj: object, seen: set | None = None) -> None:
+    """Replace every `body` string anywhere in the tree, not just jobs.*.body.
+
+    A project export nests its jobs under each workflow, so a walk that only
+    looks at the top level hands those bodies straight to the model.
+    """
+    if seen is None:
+        seen = set()
+    if id(obj) in seen:
+        return
+    if isinstance(obj, dict):
+        seen.add(id(obj))
+        for key, value in obj.items():
+            if key == "body" and isinstance(value, str):
+                obj[key] = REDACTED_BODY
+            else:
+                _redact_bodies(value, seen)
+    elif isinstance(obj, list):
+        seen.add(id(obj))
+        for item in obj:
+            _redact_bodies(item, seen)
+
+
 def redact_job_bodies(yaml_str: str) -> str:
     """Return workflow YAML with job bodies replaced by a placeholder and id
     fields removed.
 
     This is the read-only structural view shown to the planner and to job_chat
-    in subagent mode. It never round-trips back into a real workflow, so the
-    UUID ids are pure noise to the model — dropping them saves tokens.
+    in subagent mode. Bodies are deferred rather than hidden: the model reads
+    any of them with inspect_job_code, so this is about tokens, not secrecy.
+    The UUID ids never round-trip back into a real workflow, so dropping them
+    saves tokens too.
+
+    A document we cannot read is returned as it came. The model can say what is
+    wrong with it, which is more use than telling it there is no workflow.
     """
     try:
         yaml_data = yaml.safe_load(yaml_str)
-        if yaml_data and "jobs" in yaml_data:
-            _remove_ids(yaml_data)
-            for job_data in yaml_data["jobs"].values():
-                if "body" in job_data:
-                    job_data["body"] = "# [use inspect_job_code to view]"
-            return yaml.dump(yaml_data, sort_keys=False)
+        if not isinstance(yaml_data, dict):
+            return yaml_str
+        _remove_ids(yaml_data)
+        _redact_bodies(yaml_data)
+        return yaml.dump(yaml_data, sort_keys=False)
     except Exception:
-        pass
-    return yaml_str
+        return yaml_str
+
 
+def _remove_ids(obj: object, seen: set | None = None) -> None:
+    """Recursively remove 'id' keys from a parsed YAML structure.
 
-def _remove_ids(obj: object) -> None:
-    """Recursively remove 'id' keys from a parsed YAML structure."""
+    A YAML anchor can refer to its own container, and PyYAML builds that as a
+    real cycle, so the walk tracks what it has already entered.
+    """
+    if seen is None:
+        seen = set()
+    if id(obj) in seen:
+        return
     if isinstance(obj, dict):
+        seen.add(id(obj))
         obj.pop("id", None)
         for value in obj.values():
-            _remove_ids(value)
+            _remove_ids(value, seen)
     elif isinstance(obj, list):
+        seen.add(id(obj))
         for item in obj:
-            _remove_ids(item)
+            _remove_ids(item, seen)
 
 
 def stitch_job_code(yaml_str: str, job_key: str, new_code: str) -> str: