diff --git a/src/test_work_graph_plan.py b/src/test_work_graph_plan.py index 0cb23bf..aa6ae76 100644 --- a/src/test_work_graph_plan.py +++ b/src/test_work_graph_plan.py @@ -62,6 +62,51 @@ class PlanWorkGraphTestBase(unittest.TestCase): self.conn.close() +class ExtractJsonPayloadTest(unittest.TestCase): + def test_bare_json_result(self): + payload = {"content": "hello"} + self.assertEqual(work_graph_plan.extract_json_payload(json.dumps(payload), ["content"]), payload) + + def test_fenced_json_block_with_prose_around_it(self): + payload = {"content": "hello"} + text = "Here you go:\n\n```json\n" + json.dumps(payload) + "\n```\n\nDone." + self.assertEqual(work_graph_plan.extract_json_payload(text, ["content"]), payload) + + def test_missing_required_key_returns_none(self): + text = "```json\n" + json.dumps({"other": "x"}) + "\n```" + self.assertIsNone(work_graph_plan.extract_json_payload(text, ["content"])) + + def test_no_json_anywhere_returns_none(self): + self.assertIsNone(work_graph_plan.extract_json_payload("sorry, I can't do that", ["content"])) + + def test_content_containing_braces_extracts_correctly(self): + # a real code-writing node's content is guaranteed to contain '{'/'}' + # (dict literals, f-strings, function calls) -- must not confuse + # extraction just because the payload itself contains braces. + payload = {"content": 'def f():\n d = {"a": 1, "b": 2}\n return d\n'} + text = "```json\n" + json.dumps(payload) + "\n```" + self.assertEqual(work_graph_plan.extract_json_payload(text, ["content"]), payload) + + def test_content_with_brace_then_whitespace_then_backticks_does_not_truncate(self): + # Real incident (2026-08-15, task b243dfe8), diagnosed live with the + # user after they pushed back on "is this over-constrained code" -- + # the honest answer was the opposite: the OLD implementation + # approximated JSON-boundary-finding with a non-greedy regex + # (```(?:json)?\s*(\{.*?\})\s*```), which stops at the FIRST '}' + # anywhere in the text followed by whitespace + a closing fence. + # This is fully VALID, json.dumps()-produced JSON -- no malformed + # input required -- the content just happens to contain a dict + # literal followed by spaces (legal, unescaped JSON string content) + # and then a literal ``` sequence, e.g. a script whose own docstring + # shows an example. The regex silently returned None here; a real + # JSON decoder (json.JSONDecoder.raw_decode, which this function now + # uses) correctly walks past it because it understands JSON string/ + # brace grammar instead of approximating it. + payload = {"content": 'd = {"a": 1} ```\nmore real code follows this line\n'} + text = "```json\n" + json.dumps(payload) + "\n```" + self.assertEqual(work_graph_plan.extract_json_payload(text, ["content"]), payload) + + class ExtractGraphManifestTest(unittest.TestCase): def test_bare_json_result(self): self.assertEqual(work_graph_plan._extract_graph_manifest(json.dumps(_MANIFEST)), _MANIFEST) @@ -76,6 +121,20 @@ class ExtractGraphManifestTest(unittest.TestCase): def test_no_json_anywhere_returns_none(self): self.assertIsNone(work_graph_plan._extract_graph_manifest("sorry, I can't do that")) + def test_wrapper_key_content_containing_braces_does_not_truncate(self): + # Same class of bug as ExtractJsonPayloadTest's own brace/whitespace/ + # backtick case, for this function's OWN separate wrapper-key + # fallback scan (previously also regex-based, same vulnerability). + manifest_with_brace_content = { + "nodes": [{ + "node_id": "n1", "artifact_type": "brd", "depends_on": [], + "note": 'example: {"x": 1} ```\nmore text\n', + }], + "edges": [], + } + text = _fenced_result(manifest_with_brace_content, wrapper_key="graph_manifest") + self.assertEqual(work_graph_plan._extract_graph_manifest(text), manifest_with_brace_content) + class NormalizeGraphManifestTest(unittest.TestCase): """U-031 regression: a real Planner call (2026-08-01) came back using diff --git a/src/work_graph_plan.py b/src/work_graph_plan.py index f72d27f..a11b8de 100644 --- a/src/work_graph_plan.py +++ b/src/work_graph_plan.py @@ -26,7 +26,6 @@ U-009/U-010. """ import json -import re import subprocess import uuid @@ -340,8 +339,6 @@ def _generic_prompt_template(org_node): "shaped exactly like {{\"nodes\": [...], \"edges\": [...]}} -- no extra wrapper key." ) -_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL) - class PlanningFailedError(Exception): """PLANNING_FAILED -- 03-C ยง3.4.""" @@ -350,20 +347,51 @@ class PlanningFailedError(Exception): def extract_json_payload(text, required_keys=()): """--json-schema doesn't constrain output when --agent is set (see module docstring) -- pulls a JSON payload out of free-form response text - instead: try the whole text as JSON, then each ```json fenced block; - accept the first candidate that parses to a dict containing every key in - required_keys. Generalized (2026-08, work-graph-node-execution- - deblackbox) from what used to be `_extract_graph_manifest`'s own inline - algorithm, so `org_node_trigger.trigger_org_node_structured` can reuse it - for arbitrary output_schemas, not just graph_manifest's {nodes, edges} - shape. Wrapper-key fallbacks (like _extract_graph_manifest's own - "graph_manifest" key below) are caller-specific and deliberately left - out of this shared layer.""" - candidates = [text] + _JSON_FENCE_RE.findall(text) - for candidate in candidates: + instead: try the whole text as JSON, then a real JSON object starting at + every '{' in the text; accept the first candidate that parses to a dict + containing every key in required_keys. Generalized (2026-08, work-graph- + node-execution-deblackbox) from what used to be `_extract_graph_manifest`'s + own inline algorithm, so `org_node_trigger.trigger_org_node_structured` + can reuse it for arbitrary output_schemas, not just graph_manifest's + {nodes, edges} shape. Wrapper-key fallbacks (like _extract_graph_manifest's + own "graph_manifest" key below) are caller-specific and deliberately left + out of this shared layer. + + 2026-08-15 (real incident, task b243dfe8, diagnosed live with the user -- + they correctly pushed back on "is this over-constrained code" and asked + why a check step needs to be this strict; the honest answer turned out to + be the opposite: this used to under-use what's available, approximating + JSON-object-boundary-finding with a regex, `_JSON_FENCE_RE`'s + ```(?:json)?\\s*(\\{.*?\\})\\s*``` -- a NON-GREEDY match, so it stops at the + FIRST '}' anywhere in the text that happens to be followed by whitespace + and a closing fence. Content containing real code (a test script, + guaranteed to contain '{'/'}' -- dict literals, f-strings, function + calls) can easily contain that exact sequence INCIDENTALLY, deep inside + an otherwise perfectly valid JSON string value -- no malformed input + required. Confirmed with fully valid, json.dumps()-produced input: + {"content": "d = {\"a\": 1} ```\\nmore code\\n"} silently failed to + extract, because the regex stopped at the '}' right after `{"a": 1}`, + which happens to be followed by spaces (legal, unescaped JSON string + content) and then a literal ```. A regex fundamentally can't track + nested/balanced braces correctly (a classical limitation, not a bug in + how carefully this one was written) -- the fix isn't a tighter regex, it's + using the parser Python already ships: json.JSONDecoder.raw_decode() + walking real JSON grammar (proper string-escaping/nesting awareness) is + immune to this whole class of failure, for free, with less code than the + regex approach had.""" + try: + parsed = json.loads(text) + if isinstance(parsed, dict) and all(key in parsed for key in required_keys): + return parsed + except (json.JSONDecodeError, TypeError): + pass + decoder = json.JSONDecoder() + for idx, ch in enumerate(text): + if ch != "{": + continue try: - parsed = json.loads(candidate) - except (json.JSONDecodeError, TypeError): + parsed, _ = decoder.raw_decode(text, idx) + except json.JSONDecodeError: continue if isinstance(parsed, dict) and all(key in parsed for key in required_keys): return parsed @@ -375,19 +403,37 @@ def _extract_graph_manifest(text): wrapped one level deep under a "graph_manifest" key (models sometimes add that wrapper despite being asked not to) -- kept as its own function rather than folding the wrapper-key check into extract_json_payload - itself, since no other caller needs this specific wrapper-key name.""" + itself, since no other caller needs this specific wrapper-key name. + + 2026-08-15: own fallback scan uses the same real-parser approach + extract_json_payload's own docstring explains (not the old regex, + _JSON_FENCE_RE -- same non-greedy-can't-track-nesting vulnerability + would otherwise still exist here, just for the wrapper-key shape + instead of the bare {nodes, edges} shape).""" result = extract_json_payload(text, required_keys=["nodes", "edges"]) if result is not None: return result - for candidate in [text] + _JSON_FENCE_RE.findall(text): + try: + parsed = json.loads(text) + except (json.JSONDecodeError, TypeError): + parsed = None + if isinstance(parsed, dict): + inner = parsed.get("graph_manifest") + if isinstance(inner, dict) and "nodes" in inner and "edges" in inner: + return inner + decoder = json.JSONDecoder() + for idx, ch in enumerate(text): + if ch != "{": + continue try: - parsed = json.loads(candidate) - except (json.JSONDecodeError, TypeError): + parsed, _ = decoder.raw_decode(text, idx) + except json.JSONDecodeError: + continue + if not isinstance(parsed, dict): continue - if isinstance(parsed, dict): - inner = parsed.get("graph_manifest") - if isinstance(inner, dict) and "nodes" in inner and "edges" in inner: - return inner + inner = parsed.get("graph_manifest") + if isinstance(inner, dict) and "nodes" in inner and "edges" in inner: + return inner return None