fix(kg-extract): drop ReDoS-prone regex in _parse_json_blob

[jphein]

Summary

• Replace the _JSON_ARRAY_PATTERN regex (catastrophic backtracking on .*? inside (?:...)* with re.DOTALL) with _scan_balanced_array, a linear bracket-counting scanner that respects JSON string literals.
• Two new regression tests in tests/test_kg_extractor.py lock in the fix.

Root cause

The non-greedy .*? inside the repeating non-capturing group (?:,\s*\{.*?\}\s*)* made the engine retry every combination of where each } could match. On long, unbalanced responses (e.g. a model dump truncated mid-object) the worst-case is exponential in input length.

Production symptom

py-spy dump against the KG extraction worker showed the MainThread blocked in _parse_json_blob for seconds at a time, freezing the async event loop. Both extractor hosts (katana 2080 Ti and familiar P102) sat at ~0% GPU because the workers couldn't drain LLM responses fast enough to feed continuous-batching.

After the fix:

• katana GPU climbed to ~78% utilization
• familiar GPU climbed to ~43% utilization
• 28 successful HTTP 200s in 30s on katana alone

The fix

_scan_balanced_array walks the input once:

• tracks string-literal context so a [ inside a JSON string value doesn't unbalance the depth counter
• handles escape sequences inside strings
• returns the first balanced [...] substring, or None if no balanced bracket exists

O(n) time, no backtracking, no event-loop stalls.

Tests

• test_parse_json_blob_no_redos_on_pathological_input — pumps "[" + "{" * 5000 through the parser and asserts completion in under 100ms. The old regex hung indefinitely on this input.
• test_parse_json_blob_brackets_inside_strings — confirms the string-aware scanner returns the right boundary when a [ appears inside a JSON string value (e.g. an object literal containing "[draft]").

All 30 tests in tests/test_kg_extractor.py pass in 0.10s.

Test plan

• Run tests/test_kg_extractor.py locally — 30/30 pass
• Hot-deploy on the production KG extractor on familiar.jphe.in — GPU util jumped from ~0% to ~43%
• Verify no regressions in _parse_json_blob happy-path shapes (clean array, object-wrapping-triples, prose leakage)

🤖 Generated with Claude Code

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical performance issue in the KG extraction worker where a regex-based JSON parser was causing catastrophic backtracking on malformed inputs. By replacing the regex with a deterministic, linear-time bracket-counting scanner, the fix eliminates event loop freezes and significantly improves throughput in production environments.

Highlights

• Performance Improvement: Replaced a ReDoS-prone regular expression with a linear O(n) bracket-counting scanner to prevent catastrophic backtracking and event loop stalls.
• Robustness: The new scanner correctly handles JSON string literals, ensuring that brackets contained within strings do not interfere with the balancing logic.
• Regression Testing: Added two new tests to verify performance on pathological inputs and ensure correct parsing when brackets appear inside string values.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request replaces a regex-based JSON array extraction with a linear, bracket-counting scanner (_scan_balanced_array) to prevent catastrophic backtracking and event loop freezes on pathological inputs. The feedback suggests enhancing this scanner and the parser to support a start_pos parameter and a retry loop. This ensures that if the LLM output contains leading bracketed prose (e.g., [draft]) before the actual JSON array, the parser can skip the invalid candidate and successfully extract the valid JSON array.

[gemini-code-assist]

If the LLM output contains a false start or bracketed prose (e.g., Here is the [draft] JSON: [...]), the current implementation will extract the first balanced bracketed substring ([draft]), fail to parse it as JSON, and immediately return [], ignoring the actual JSON array that follows. Implementing a loop to try subsequent [ candidates ensures the parser successfully extracts the valid JSON array even in the presence of leading bracketed prose.

Suggested change

	candidate = _scan_balanced_array(raw_stripped)
	if candidate is None:
	logger.debug("no JSON array found in LLM response: %r", raw_stripped[:200])
	return []
	try:
	parsed = json.loads(match.group(0))
	parsed = json.loads(candidate)
	except (json.JSONDecodeError, ValueError) as e:
	logger.debug("regex-extracted JSON failed to parse: %s", e)
	logger.debug("scanned JSON failed to parse: %s", e)
	return []
	start_pos = 0
	while True:
	candidate = _scan_balanced_array(raw_stripped, start_pos)
	if candidate is None:
	logger.debug("no JSON array found in LLM response: %r", raw_stripped[:200])
	return []
	try:
	parsed = json.loads(candidate)
	break
	except (json.JSONDecodeError, ValueError) as e:
	logger.debug("scanned JSON failed to parse: %s", e)
	current_start = raw_stripped.find("[", start_pos)
	start_pos = current_start + 1

[gemini-code-assist]

To support robust parsing when there are multiple bracketed expressions (e.g., false starts or bracketed prose like [draft] before the actual JSON array), update _scan_balanced_array to accept a start_pos parameter. This allows the parser to scan for subsequent candidates if the first one fails to parse as valid JSON.

Suggested change

	def _scan_balanced_array(s: str) -> Optional[str]:
	"""Return the first balanced ``[...]`` substring, respecting JSON strings.
	Linear scan, no regex backtracking. Used when the model leaks prose
	around its JSON output and a strict ``json.loads`` fails. The previous
	regex (``\\[\\s*\\{.*?\\}...\\]`` with re.DOTALL) had catastrophic
	backtracking on long, malformed inputs and could freeze the async
	event loop for seconds at a time.
	"""
	start = s.find("[")
	if start < 0:
	return None
	def _scan_balanced_array(s: str, start_pos: int = 0) -> Optional[str]:
	"""Return the first balanced ``[...]`` substring, respecting JSON strings.
	Linear scan, no regex backtracking. Used when the model leaks prose
	around its JSON output and a strict ``json.loads`` fails. The previous
	regex (``\\[\\s*\\{.*?\\}\\s*(?:,\\s*\\{.*?\\}\\s*)*\\]`` with re.DOTALL) had catastrophic
	backtracking on long, malformed inputs and could freeze the async
	event loop for seconds at a time.
	"""
	start = s.find("[", start_pos)
	if start < 0:
	return None

[Copilot]

Pull request overview

This PR hardens KG LLM response parsing by removing a ReDoS-prone regex from _parse_json_blob and replacing it with a linear-time bracket-counting scanner, with new regression tests to prevent performance regressions.

Changes:

• Replaced _JSON_ARRAY_PATTERN regex extraction with _scan_balanced_array to avoid catastrophic backtracking on malformed inputs.
• Updated _parse_json_blob to use the new scanner and adjusted debug messaging accordingly.
• Added two regression tests covering pathological input performance and bracket handling inside JSON strings.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
mempalace/kg_llm_extractor.py	Removes ReDoS-prone regex parsing and introduces a linear balanced-array scanner used by _parse_json_blob.
tests/test_kg_extractor.py	Adds regression tests for the new parsing behavior and prior performance failure mode.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.