宜博 / lemongraph

Commit 98d86e473d86d52096d93012b652c010f193e177

Authored by bboonnee
1 parent 1402836d

fix: extract_json_payload uses a real JSON decoder instead of a regex approximation 

Real incident (2026-08-15, task b243dfe8), diagnosed live with the user after
they 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 of "too many rules": the old implementation UNDER-used what's
available, approximating JSON-object-boundary-finding with a regex
(_JSON_FENCE_RE, ```(?:json)?\s*(\{.*?\})\s*```) instead of using the real
parser Python already ships.

That regex is non-greedy, so it stops at the FIRST '}' anywhere in the
response 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 trivially 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}   ```\n..."}
silently failed to extract, because the regex stopped right after `{"a": 1}`
(legal spaces, then a literal ```), discarding the real, complete payload
and raising StructuredOutputExtractionError on genuinely good content. This
is exactly what happened to a real Coder node writing a real test script
today.

A regex fundamentally can't track nested/balanced braces correctly (a
classical limitation, not a carefulness gap) -- fixed by using
json.JSONDecoder.raw_decode() to attempt a real decode starting at every
'{' in the text, which understands JSON's actual string-escaping/nesting
grammar and is immune to this whole class of failure. Same fix applied to
_extract_graph_manifest's own separate wrapper-key fallback scan (same
regex, same vulnerability, different key shape). _JSON_FENCE_RE and its
now-unused `import re` are removed -- less code, not more, and a strictly
more correct mechanism.

New ExtractJsonPayloadTest class (extract_json_payload had no direct unit
tests before this) plus a wrapper-key-scan test on
ExtractGraphManifestTest, both driving the exact valid-JSON-with-braces
scenario that broke the old regex. 502/502 across every file this and
today's other fixes touch.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Inline Side-by-side
Showing 2 changed files with 129 additions and 24 deletions
src/test_work_graph_plan.py
  View file @98d86e4
... ... @@ -62,6 +62,51 @@ class PlanWorkGraphTestBase(unittest.TestCase):
62 62 self.conn.close()
63 63  
64 64  
  65 +class ExtractJsonPayloadTest(unittest.TestCase):
  66 + def test_bare_json_result(self):
  67 + payload = {"content": "hello"}
  68 + self.assertEqual(work_graph_plan.extract_json_payload(json.dumps(payload), ["content"]), payload)
  69 +
  70 + def test_fenced_json_block_with_prose_around_it(self):
  71 + payload = {"content": "hello"}
  72 + text = "Here you go:\n\n```json\n" + json.dumps(payload) + "\n```\n\nDone."
  73 + self.assertEqual(work_graph_plan.extract_json_payload(text, ["content"]), payload)
  74 +
  75 + def test_missing_required_key_returns_none(self):
  76 + text = "```json\n" + json.dumps({"other": "x"}) + "\n```"
  77 + self.assertIsNone(work_graph_plan.extract_json_payload(text, ["content"]))
  78 +
  79 + def test_no_json_anywhere_returns_none(self):
  80 + self.assertIsNone(work_graph_plan.extract_json_payload("sorry, I can't do that", ["content"]))
  81 +
  82 + def test_content_containing_braces_extracts_correctly(self):
  83 + # a real code-writing node's content is guaranteed to contain '{'/'}'
  84 + # (dict literals, f-strings, function calls) -- must not confuse
  85 + # extraction just because the payload itself contains braces.
  86 + payload = {"content": 'def f():\n d = {"a": 1, "b": 2}\n return d\n'}
  87 + text = "```json\n" + json.dumps(payload) + "\n```"
  88 + self.assertEqual(work_graph_plan.extract_json_payload(text, ["content"]), payload)
  89 +
  90 + def test_content_with_brace_then_whitespace_then_backticks_does_not_truncate(self):
  91 + # Real incident (2026-08-15, task b243dfe8), diagnosed live with the
  92 + # user after they pushed back on "is this over-constrained code" --
  93 + # the honest answer was the opposite: the OLD implementation
  94 + # approximated JSON-boundary-finding with a non-greedy regex
  95 + # (```(?:json)?\s*(\{.*?\})\s*```), which stops at the FIRST '}'
  96 + # anywhere in the text followed by whitespace + a closing fence.
  97 + # This is fully VALID, json.dumps()-produced JSON -- no malformed
  98 + # input required -- the content just happens to contain a dict
  99 + # literal followed by spaces (legal, unescaped JSON string content)
  100 + # and then a literal ``` sequence, e.g. a script whose own docstring
  101 + # shows an example. The regex silently returned None here; a real
  102 + # JSON decoder (json.JSONDecoder.raw_decode, which this function now
  103 + # uses) correctly walks past it because it understands JSON string/
  104 + # brace grammar instead of approximating it.
  105 + payload = {"content": 'd = {"a": 1} ```\nmore real code follows this line\n'}
  106 + text = "```json\n" + json.dumps(payload) + "\n```"
  107 + self.assertEqual(work_graph_plan.extract_json_payload(text, ["content"]), payload)
  108 +
  109 +
65 110 class ExtractGraphManifestTest(unittest.TestCase):
66 111 def test_bare_json_result(self):
67 112 self.assertEqual(work_graph_plan._extract_graph_manifest(json.dumps(_MANIFEST)), _MANIFEST)
... ... @@ -76,6 +121,20 @@ class ExtractGraphManifestTest(unittest.TestCase):
76 121 def test_no_json_anywhere_returns_none(self):
77 122 self.assertIsNone(work_graph_plan._extract_graph_manifest("sorry, I can't do that"))
78 123  
  124 + def test_wrapper_key_content_containing_braces_does_not_truncate(self):
  125 + # Same class of bug as ExtractJsonPayloadTest's own brace/whitespace/
  126 + # backtick case, for this function's OWN separate wrapper-key
  127 + # fallback scan (previously also regex-based, same vulnerability).
  128 + manifest_with_brace_content = {
  129 + "nodes": [{
  130 + "node_id": "n1", "artifact_type": "brd", "depends_on": [],
  131 + "note": 'example: {"x": 1} ```\nmore text\n',
  132 + }],
  133 + "edges": [],
  134 + }
  135 + text = _fenced_result(manifest_with_brace_content, wrapper_key="graph_manifest")
  136 + self.assertEqual(work_graph_plan._extract_graph_manifest(text), manifest_with_brace_content)
  137 +
79 138  
80 139 class NormalizeGraphManifestTest(unittest.TestCase):
81 140 """U-031 regression: a real Planner call (2026-08-01) came back using
... ...
src/work_graph_plan.py
  View file @98d86e4
... ... @@ -26,7 +26,6 @@ U-009/U-010.
26 26 """
27 27  
28 28 import json
29   -import re
30 29 import subprocess
31 30 import uuid
32 31  
... ... @@ -340,8 +339,6 @@ def _generic_prompt_template(org_node):
340 339 "shaped exactly like {{\"nodes\": [...], \"edges\": [...]}} -- no extra wrapper key."
341 340 )
342 341  
343   -_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL)
344   -
345 342  
346 343 class PlanningFailedError(Exception):
347 344 """PLANNING_FAILED -- 03-C §3.4."""
... ... @@ -350,20 +347,51 @@ class PlanningFailedError(Exception):
350 347 def extract_json_payload(text, required_keys=()):
351 348 """--json-schema doesn't constrain output when --agent is set (see module
352 349 docstring) -- pulls a JSON payload out of free-form response text
353   - instead: try the whole text as JSON, then each ```json fenced block;
354   - accept the first candidate that parses to a dict containing every key in
355   - required_keys. Generalized (2026-08, work-graph-node-execution-
356   - deblackbox) from what used to be `_extract_graph_manifest`'s own inline
357   - algorithm, so `org_node_trigger.trigger_org_node_structured` can reuse it
358   - for arbitrary output_schemas, not just graph_manifest's {nodes, edges}
359   - shape. Wrapper-key fallbacks (like _extract_graph_manifest's own
360   - "graph_manifest" key below) are caller-specific and deliberately left
361   - out of this shared layer."""
362   - candidates = [text] + _JSON_FENCE_RE.findall(text)
363   - for candidate in candidates:
  350 + instead: try the whole text as JSON, then a real JSON object starting at
  351 + every '{' in the text; accept the first candidate that parses to a dict
  352 + containing every key in required_keys. Generalized (2026-08, work-graph-
  353 + node-execution-deblackbox) from what used to be `_extract_graph_manifest`'s
  354 + own inline algorithm, so `org_node_trigger.trigger_org_node_structured`
  355 + can reuse it for arbitrary output_schemas, not just graph_manifest's
  356 + {nodes, edges} shape. Wrapper-key fallbacks (like _extract_graph_manifest's
  357 + own "graph_manifest" key below) are caller-specific and deliberately left
  358 + out of this shared layer.
  359 +
  360 + 2026-08-15 (real incident, task b243dfe8, diagnosed live with the user --
  361 + they correctly pushed back on "is this over-constrained code" and asked
  362 + why a check step needs to be this strict; the honest answer turned out to
  363 + be the opposite: this used to under-use what's available, approximating
  364 + JSON-object-boundary-finding with a regex, `_JSON_FENCE_RE`'s
  365 + ```(?:json)?\\s*(\\{.*?\\})\\s*``` -- a NON-GREEDY match, so it stops at the
  366 + FIRST '}' anywhere in the text that happens to be followed by whitespace
  367 + and a closing fence. Content containing real code (a test script,
  368 + guaranteed to contain '{'/'}' -- dict literals, f-strings, function
  369 + calls) can easily contain that exact sequence INCIDENTALLY, deep inside
  370 + an otherwise perfectly valid JSON string value -- no malformed input
  371 + required. Confirmed with fully valid, json.dumps()-produced input:
  372 + {"content": "d = {\"a\": 1} ```\\nmore code\\n"} silently failed to
  373 + extract, because the regex stopped at the '}' right after `{"a": 1}`,
  374 + which happens to be followed by spaces (legal, unescaped JSON string
  375 + content) and then a literal ```. A regex fundamentally can't track
  376 + nested/balanced braces correctly (a classical limitation, not a bug in
  377 + how carefully this one was written) -- the fix isn't a tighter regex, it's
  378 + using the parser Python already ships: json.JSONDecoder.raw_decode()
  379 + walking real JSON grammar (proper string-escaping/nesting awareness) is
  380 + immune to this whole class of failure, for free, with less code than the
  381 + regex approach had."""
  382 + try:
  383 + parsed = json.loads(text)
  384 + if isinstance(parsed, dict) and all(key in parsed for key in required_keys):
  385 + return parsed
  386 + except (json.JSONDecodeError, TypeError):
  387 + pass
  388 + decoder = json.JSONDecoder()
  389 + for idx, ch in enumerate(text):
  390 + if ch != "{":
  391 + continue
364 392 try:
365   - parsed = json.loads(candidate)
366   - except (json.JSONDecodeError, TypeError):
  393 + parsed, _ = decoder.raw_decode(text, idx)
  394 + except json.JSONDecodeError:
367 395 continue
368 396 if isinstance(parsed, dict) and all(key in parsed for key in required_keys):
369 397 return parsed
... ... @@ -375,19 +403,37 @@ def _extract_graph_manifest(text):
375 403 wrapped one level deep under a "graph_manifest" key (models sometimes
376 404 add that wrapper despite being asked not to) -- kept as its own function
377 405 rather than folding the wrapper-key check into extract_json_payload
378   - itself, since no other caller needs this specific wrapper-key name."""
  406 + itself, since no other caller needs this specific wrapper-key name.
  407 +
  408 + 2026-08-15: own fallback scan uses the same real-parser approach
  409 + extract_json_payload's own docstring explains (not the old regex,
  410 + _JSON_FENCE_RE -- same non-greedy-can't-track-nesting vulnerability
  411 + would otherwise still exist here, just for the wrapper-key shape
  412 + instead of the bare {nodes, edges} shape)."""
379 413 result = extract_json_payload(text, required_keys=["nodes", "edges"])
380 414 if result is not None:
381 415 return result
382   - for candidate in [text] + _JSON_FENCE_RE.findall(text):
  416 + try:
  417 + parsed = json.loads(text)
  418 + except (json.JSONDecodeError, TypeError):
  419 + parsed = None
  420 + if isinstance(parsed, dict):
  421 + inner = parsed.get("graph_manifest")
  422 + if isinstance(inner, dict) and "nodes" in inner and "edges" in inner:
  423 + return inner
  424 + decoder = json.JSONDecoder()
  425 + for idx, ch in enumerate(text):
  426 + if ch != "{":
  427 + continue
383 428 try:
384   - parsed = json.loads(candidate)
385   - except (json.JSONDecodeError, TypeError):
  429 + parsed, _ = decoder.raw_decode(text, idx)
  430 + except json.JSONDecodeError:
  431 + continue
  432 + if not isinstance(parsed, dict):
386 433 continue
387   - if isinstance(parsed, dict):
388   - inner = parsed.get("graph_manifest")
389   - if isinstance(inner, dict) and "nodes" in inner and "edges" in inner:
390   - return inner
  434 + inner = parsed.get("graph_manifest")
  435 + if isinstance(inner, dict) and "nodes" in inner and "edges" in inner:
  436 + return inner
391 437 return None
392 438  
393 439  
... ...