from __future__ import annotations import ast import difflib import itertools import json import re import statistics from collections import Counter from pathlib import Path from typing import Any, Dict, List import libcst as cst try: from ..models import ASTSummary, FailurePattern, RunRecord except ImportError: from models import ASTSummary, FailurePattern, RunRecord def list_repo_structure(root_path: str) -> List[dict]: root = Path(root_path) ignored_dirs = {"__pycache__", ".git", "node_modules", "venv", ".venv"} entries: List[dict] = [] for path in root.rglob("*.py"): if any(part in ignored_dirs for part in path.parts): continue content = path.read_text(encoding="utf-8", errors="ignore") name = path.name entries.append( { "path": str(path.relative_to(root)).replace("\\", "/"), "is_test": name.startswith("test_") or name.endswith("_test.py"), "size_bytes": path.stat().st_size, "has_async": ("async def" in content) or ("await " in content), } ) return entries def read_file_excerpt(path: str, start_line: int, end_line: int) -> str: if end_line - start_line > 100: raise ValueError("Cannot read more than 100 lines per call") p = Path(path) with p.open("r", encoding="utf-8", errors="ignore") as f: lines = itertools.islice(f, max(start_line - 1, 0), end_line) return "".join(lines) def parse_ast_summary(path: str) -> ASTSummary: p = Path(path) if p.suffix in {".js", ".ts"}: return _parse_js_ts_summary(p) source = p.read_text(encoding="utf-8", errors="ignore") return _parse_python_ast_summary(source) def resolve_target_from_evidence(path: str, evidence: List[str]) -> Dict[str, Any]: summary = parse_ast_summary(path) if not evidence: if summary.functions: return { "type": "function", "identifier": summary.functions[0]["name"], "line": summary.functions[0]["line_range"][0], } return {"type": "line", "identifier": "def test_", "line": 1} source = Path(path).read_text(encoding="utf-8", errors="ignore") for clue in evidence: clue = clue.strip() if not clue: continue fn_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*)", clue) if fn_match: fn_name = fn_match.group(1) for fn in summary.functions: if fn["name"] == fn_name: return { "type": "function", "identifier": fn_name, "line": fn["line_range"][0], } if clue in source: return {"type": "line", "identifier": clue, "line": _line_of_text(source, clue)} if summary.external_calls: return {"type": "call", "identifier": summary.external_calls[0], "line": 1} if summary.functions: return {"type": "function", "identifier": summary.functions[0]["name"], "line": summary.functions[0]["line_range"][0]} return {"type": "line", "identifier": "def test_", "line": 1} class _LoggingInjector(cst.CSTTransformer): def __init__(self, targets: List[Dict[str, str]]) -> None: self.targets = targets def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: target = next((t for t in self.targets if t.get("function_name") == original_node.name.value), None) if not target: return updated_node position = target.get("position", "entry") fn_name = original_node.name.value log_stmt = cst.parse_statement( "print(json.dumps({" "'timestamp': time.time(), " "'thread_id': threading.get_ident(), " "'task_id': _task_name(), " f"'function': '{fn_name}', " f"'event': '{position}'" "}), file=sys.stderr)\n" ) body_stmts = list(updated_node.body.body) already_present = any("'function': '%s'" % fn_name in getattr(stmt, "code", "") for stmt in body_stmts) if already_present: return updated_node if position == "exit": body_stmts.append(log_stmt) else: body_stmts.insert(0, log_stmt) return updated_node.with_changes(body=updated_node.body.with_changes(body=body_stmts)) def inject_logging(path: str, injection_points: List[Dict[str, str]]) -> str: source = Path(path).read_text(encoding="utf-8", errors="ignore") module = cst.parse_module(source) header = ( "import asyncio\n" "import json\n" "import sys\n" "import threading\n" "import time\n\n" "def _task_name():\n" " try:\n" " task = asyncio.current_task()\n" " return task.get_name() if task else None\n" " except Exception:\n" " return None\n\n" ) transformed = module.visit(_LoggingInjector(injection_points)).code if "def _task_name():" not in transformed: transformed = header + transformed return transformed def apply_ast_patch(path: str, patch_spec: Dict[str, Any]) -> Dict[str, Any]: p = Path(path) original = p.read_text(encoding="utf-8", errors="ignore") try: operation = patch_spec.get("operation") target = patch_spec.get("target", {}) identifier = str(target.get("identifier", "")) template = str(patch_spec.get("code_template", "")) params = patch_spec.get("parameters", {}) rendered = template.format(**params) if rendered and rendered in original and operation in {"insert_before", "insert_after", "add_decorator"}: diff_data = compute_diff_from_sources(str(p), original, original) return { "success": True, "diff": diff_data["unified_diff"], "lines_changed": 0, "ast_diff": diff_data["ast_diff"], } updated = _apply_patch_operation(original, operation, target, identifier, rendered) p.write_text(updated, encoding="utf-8") diff_data = compute_diff_from_sources(str(p), original, updated) return { "success": True, "diff": diff_data["unified_diff"], "lines_changed": diff_data["lines_changed"], "ast_diff": diff_data["ast_diff"], } except Exception as exc: return { "success": False, "diff": "", "lines_changed": 0, "ast_diff": {}, "error": str(exc), } def compute_diff(original_path: str, patched_source: str) -> Dict[str, Any]: original = Path(original_path).read_text(encoding="utf-8", errors="ignore") return compute_diff_from_sources(original_path, original, patched_source) def compute_diff_from_sources(path: str, before_source: str, after_source: str) -> Dict[str, Any]: unified_diff = _make_unified_diff(before_source, after_source, path) lines_changed = _count_diff_lines(unified_diff) before = _safe_summary_from_source(before_source) after = _safe_summary_from_source(after_source) before_funcs = {f["name"] for f in before.functions} after_funcs = {f["name"] for f in after.functions} before_calls = set(before.external_calls) after_calls = set(after.external_calls) ast_diff = { "functions_added": sorted(after_funcs - before_funcs), "functions_removed": sorted(before_funcs - after_funcs), "functions_modified": sorted(before_funcs & after_funcs), "calls_added": sorted(after_calls - before_calls), "calls_removed": sorted(before_calls - after_calls), } return { "unified_diff": unified_diff, "lines_changed": lines_changed, "ast_diff": ast_diff, } def get_failure_pattern(run_records: List[RunRecord]) -> FailurePattern: if not run_records: return FailurePattern( pass_rate=0.0, most_common_error=None, error_distribution={}, duration_mean=0.0, duration_std=0.0, flakiness_score=0.0, ) pass_rate = sum(1 for r in run_records if r.passed) / len(run_records) errors = [r.error_type for r in run_records if r.error_type] error_distribution = dict(Counter(errors)) most_common_error = None if error_distribution: most_common_error = max(error_distribution, key=error_distribution.get) durations = [r.duration_ms for r in run_records] duration_mean = statistics.fmean(durations) duration_std = statistics.pstdev(durations) if len(durations) > 1 else 0.0 flakiness_score = 4.0 * pass_rate * (1.0 - pass_rate) return FailurePattern( pass_rate=pass_rate, most_common_error=most_common_error, error_distribution=error_distribution, duration_mean=duration_mean, duration_std=duration_std, flakiness_score=flakiness_score, ) def get_similar_fixes(root_cause_category: str, test_source: str, embedding_model: Any = None) -> List[Dict[str, str]]: try: import chromadb from sentence_transformers import SentenceTransformer client = chromadb.PersistentClient(path=str(Path.cwd() / ".flakeforge_chroma")) collection = client.get_or_create_collection("flake_fixes") if collection.count() == 0: _seed_collection_from_seed_repos(collection) model = embedding_model or SentenceTransformer("all-MiniLM-L6-v2") query_embedding = model.encode(test_source).tolist() where = {"root_cause_category": root_cause_category} results = collection.query(query_embeddings=[query_embedding], n_results=3, where=where) docs = results.get("documents", [[]])[0] metadatas = results.get("metadatas", [[]])[0] out: List[Dict[str, str]] = [] for idx, doc in enumerate(docs): md = metadatas[idx] if idx < len(metadatas) else {} payload = json.loads(doc) out.append( { "root_cause": md.get("root_cause_category", root_cause_category), "original": payload.get("original", ""), "fixed": payload.get("fixed", ""), "diff": payload.get("diff", ""), "action": md.get("action", ""), "repo": md.get("repo", "unknown"), } ) return out[:3] except Exception: return [] def _stack_trace_file_candidates(stack_trace: str, repo_path: Path) -> List[str]: """Extract repo-relative python file paths from a stack trace.""" if not stack_trace.strip(): return [] candidates: List[str] = [] for line in stack_trace.splitlines(): m = re.search(r'File\s+"([^"]+\.py)"', line) if not m: continue raw_path = m.group(1).replace("\\", "/") # Accept both absolute and relative forms. abs_repo = repo_path.resolve().as_posix().rstrip("/") if raw_path.startswith(abs_repo): rel = raw_path[len(abs_repo):].lstrip("/") if rel: candidates.append(rel) continue if raw_path.startswith("./"): candidates.append(raw_path[2:]) else: candidates.append(raw_path) # Keep order, de-duplicate. seen = set() ordered: List[str] = [] for c in candidates: key = c.strip().replace("\\", "/") if not key or key in seen: continue seen.add(key) ordered.append(key) return ordered def build_agent_targeting_hints( repo_path: str, test_identifier: str, failing_stack_trace: str, source_under_test: str, causal_frontier: str, deep_signals: Dict[str, Any], max_hints: int = 8, ) -> List[str]: """Build file-targeting hints for the agent from multiple evidence channels. This function intentionally combines more than causal graph data: - stack trace file extraction - test identifier path - import-based source-under-test extraction - causal frontier file - deep flakiness signal file lists """ root = Path(repo_path) hints: List[str] = [] def add(path: str, reason: str, score: float) -> None: p = path.strip().replace("\\", "/") if not p: return hints.append(f"{score:.2f} | {p} | {reason}") # 1) Test file gets high priority because many flaky fixes live in tests. test_file = test_identifier.split("::", 1)[0].strip() if test_file: add(test_file, "target test from test_identifier", 1.00) # 2) Stack trace frames often provide the strongest location hints. for idx, rel in enumerate(_stack_trace_file_candidates(failing_stack_trace, root)): score = max(0.70 - idx * 0.08, 0.35) add(rel, "stack trace frame", score) # 3) Causal frontier can narrow to the deepest failure frame. if causal_frontier: frontier_file = causal_frontier.split(":", 1)[0].strip() if frontier_file: add(frontier_file, "causal frontier failure site", 0.85) # 4) Imported modules from source_under_test snapshot. if source_under_test: try: summary = _parse_python_ast_summary(source_under_test) for imp in summary.imports[:8]: module_path = imp.split(".")[0].strip() if module_path: add(f"{module_path}.py", "import in source_under_test", 0.45) except Exception: pass # 5) Deep signal file lists. signal_lists = [ ("module_cache_violations", "module-cache risk"), ("fixture_scope_risks", "fixture scope risk"), ("mock_residue_sites", "mock residue site"), ("import_side_effect_files", "import side-effect file"), ] for key, label in signal_lists: values = deep_signals.get(key) or [] if not isinstance(values, list): continue for rel in values[:5]: add(str(rel), label, 0.55) # Aggregate by path, keep max score and merged reasons. by_path: Dict[str, Dict[str, Any]] = {} for row in hints: try: score_text, path_text, reason_text = [p.strip() for p in row.split("|", 2)] score = float(score_text) except Exception: continue entry = by_path.setdefault(path_text, {"score": score, "reasons": []}) entry["score"] = max(entry["score"], score) if reason_text not in entry["reasons"]: entry["reasons"].append(reason_text) ranked = sorted( by_path.items(), key=lambda item: item[1]["score"], reverse=True, ) output: List[str] = [] for path_text, payload in ranked[:max_hints]: reasons = ", ".join(payload["reasons"][:2]) output.append(f"{path_text} (score={payload['score']:.2f}; {reasons})") return output def _seed_collection_from_seed_repos(collection: Any) -> None: root = Path.cwd() / "seed_repos" if not root.exists(): return model = None try: from sentence_transformers import SentenceTransformer model = SentenceTransformer("all-MiniLM-L6-v2") except Exception: pass mapping = { "timing_race": "TIMING_RACE", "shared_state": "SHARED_STATE", "external_dependency": "EXTERNAL_DEPENDENCY", "order_dependency": "ORDER_DEPENDENCY", "resource_leak": "RESOURCE_LEAK", "nondeterminism": "NONDETERMINISM", "compound_timing_shared": "TIMING_RACE", "compound_external_nondeterminism": "EXTERNAL_DEPENDENCY", } action_mapping = { "timing_race": "ADD_TIMING_GUARD", "shared_state": "RESET_STATE", "external_dependency": "MOCK_DEPENDENCY", "order_dependency": "ADD_SYNCHRONIZATION", "resource_leak": "RESET_STATE", "nondeterminism": "SEED_RANDOMNESS", "compound_timing_shared": "ADD_TIMING_GUARD", "compound_external_nondeterminism": "MOCK_DEPENDENCY", } docs = [] embeddings = [] ids = [] metadatas = [] for repo_dir in sorted([p for p in root.iterdir() if p.is_dir()]): test_file = repo_dir / "tests" / "test_flaky.py" fixed_test = repo_dir / "solution" / "tests" / "test_flaky.py" if not test_file.exists() or not fixed_test.exists(): continue original = test_file.read_text(encoding="utf-8", errors="ignore") fixed = fixed_test.read_text(encoding="utf-8", errors="ignore") diff = _make_unified_diff(original, fixed, str(test_file)) payload = json.dumps({"original": original, "fixed": fixed, "diff": diff}) docs.append(payload) ids.append(repo_dir.name) metadatas.append( { "repo": repo_dir.name, "root_cause_category": mapping.get(repo_dir.name, "NONDETERMINISM"), "action": action_mapping.get(repo_dir.name, "GATHER_EVIDENCE"), } ) if model is not None: embeddings.append(model.encode(original).tolist()) if not docs: return if embeddings: collection.add(ids=ids, documents=docs, metadatas=metadatas, embeddings=embeddings) else: collection.add(ids=ids, documents=docs, metadatas=metadatas) def _safe_summary_from_source(source: str) -> ASTSummary: try: return _parse_python_ast_summary(source) except Exception: return ASTSummary([], [], [], [], [], []) def _parse_python_ast_summary(source: str) -> ASTSummary: tree = ast.parse(source) functions: List[Dict[str, Any]] = [] classes: List[Dict[str, Any]] = [] imports: List[str] = [] global_vars: List[str] = [] threading_primitives: List[str] = [] external_calls: List[str] = [] external_modules = { "requests", "boto3", "httpx", "redis", "psycopg2", # V2 expansion: missing real-world dependencies "aiohttp", "grpc", "sqlalchemy", "asyncpg", "celery", "pika", "subprocess", "time", } primitive_markers = { "threading.Lock", "threading.Event", "asyncio.Event", "asyncio.Lock", "asyncio.Barrier", "asyncio.sleep", } for node in ast.walk(tree): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): functions.append( { "name": node.name, "is_async": isinstance(node, ast.AsyncFunctionDef), "args": [arg.arg for arg in node.args.args], "decorators": [ast.unparse(d) for d in node.decorator_list], "line_range": [getattr(node, "lineno", 0), getattr(node, "end_lineno", 0)], } ) elif isinstance(node, ast.ClassDef): classes.append( { "name": node.name, "bases": [ast.unparse(b) for b in node.bases], "decorators": [ast.unparse(d) for d in node.decorator_list], "line_range": [getattr(node, "lineno", 0), getattr(node, "end_lineno", 0)], } ) elif isinstance(node, ast.Import): for alias in node.names: imports.append(alias.name) elif isinstance(node, ast.ImportFrom): module = node.module or "" for alias in node.names: imports.append(f"{module}.{alias.name}" if module else alias.name) for node in tree.body: if isinstance(node, ast.Assign): for target in node.targets: if isinstance(target, ast.Name): global_vars.append(target.id) for node in ast.walk(tree): if isinstance(node, ast.Call): call_name = _node_to_name(node.func) if call_name in primitive_markers: threading_primitives.append(call_name) for module in external_modules: if call_name.startswith(f"{module}."): external_calls.append(call_name) return ASTSummary( functions=functions, classes=classes, imports=sorted(set(imports)), global_vars=sorted(set(global_vars)), threading_primitives=sorted(set(threading_primitives)), external_calls=sorted(set(external_calls)), ) def _make_unified_diff(before: str, after: str, path: str) -> str: return "".join( difflib.unified_diff( before.splitlines(keepends=True), after.splitlines(keepends=True), fromfile=f"a/{path}", tofile=f"b/{path}", n=3, ) ) def _count_diff_lines(diff: str) -> int: count = 0 for line in diff.splitlines(): if line.startswith("+++") or line.startswith("---"): continue if line.startswith("+") or line.startswith("-"): count += 1 return count def _node_to_name(node: ast.AST) -> str: if isinstance(node, ast.Name): return node.id if isinstance(node, ast.Attribute): root = _node_to_name(node.value) return f"{root}.{node.attr}" if root else node.attr return "" def _line_of_text(source: str, text: str) -> int: idx = source.find(text) if idx < 0: return 1 return source[:idx].count("\n") + 1 def _apply_patch_operation(source: str, operation: str, target: Dict[str, Any], identifier: str, rendered: str) -> str: if operation == "ensure_reset_fixture": scope = str(target.get("scope", "function")) return _apply_reset_fixture(source, scope) if operation == "ensure_retry_wrapper": max_attempts = int(target.get("max_attempts", 2)) backoff_ms = int(target.get("backoff_ms", 100)) fn_name = str(target.get("function_name", "test")) return _apply_retry_wrapper(source, fn_name, max_attempts, backoff_ms) if operation == "ensure_seed_call": library = str(target.get("library", "random")) fn_name = str(target.get("function_name", "test")) return _apply_seed_call(source, fn_name, library) if operation == "add_decorator": return _apply_add_decorator(source, target, rendered) if operation == "replace_call": return _apply_replace_call(source, identifier, rendered) if operation == "wrap_with": return _apply_wrap_with(source, target, rendered) if operation in {"insert_before", "insert_after"}: if target.get("type") == "function": return _apply_insert_in_function(source, target, rendered, operation) return _apply_textual_operation_idempotent(source, operation, identifier, rendered) # ── V2 Deep-Action Operations (Bug 2 fix) ────────────────────────────────── if operation == "refactor_concurrency_primitive": return _apply_refactor_concurrency(source, target) if operation == "isolate_boundary_call": return _apply_isolate_boundary(source, target) if operation == "extract_async_scope": return _apply_extract_async_scope(source, target) if operation == "harden_idempotency": return _apply_harden_idempotency(source, target) raise ValueError(f"Unsupported operation: {operation}") def _apply_add_decorator(source: str, target: Dict[str, Any], decorator_code: str) -> str: fn_name = str(target.get("identifier", "")).replace("def ", "").strip() if decorator_code.startswith("@"): decorator_code = decorator_code[1:] module = cst.parse_module(source) decorator_expr = cst.parse_expression(decorator_code) class DecoratorAdder(cst.CSTTransformer): def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: if fn_name and fn_name not in original_node.name.value and fn_name != "test": return updated_node existing = [d.decorator.code for d in updated_node.decorators] if decorator_expr.code in existing: return updated_node return updated_node.with_changes(decorators=[*updated_node.decorators, cst.Decorator(decorator_expr)]) return module.visit(DecoratorAdder()).code def _apply_replace_call(source: str, call_identifier: str, rendered: str) -> str: module = cst.parse_module(source) rendered_expr = cst.parse_expression(rendered) class CallReplacer(cst.CSTTransformer): def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression: name = _cst_name(updated_node.func) if name == call_identifier: return rendered_expr return updated_node return module.visit(CallReplacer()).code def _apply_wrap_with(source: str, target: Dict[str, Any], rendered: str) -> str: fn_name = str(target.get("identifier", "")).strip() with_expr = rendered.strip().split(":", 1)[0].replace("with ", "").strip() module = cst.parse_module(source) class WrapWithTransformer(cst.CSTTransformer): def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: if fn_name and fn_name not in original_node.name.value and fn_name != "test": return updated_node with_stmt = cst.With( items=[cst.WithItem(item=cst.parse_expression(with_expr))], body=cst.IndentedBlock(body=list(updated_node.body.body)), ) wrapped = with_stmt return updated_node.with_changes(body=updated_node.body.with_changes(body=[wrapped])) return module.visit(WrapWithTransformer()).code def _apply_insert_in_function(source: str, target: Dict[str, Any], rendered: str, operation: str) -> str: fn_name = str(target.get("identifier", "")).strip() statement = cst.parse_statement(rendered + "\n") module = cst.parse_module(source) class InsertTransformer(cst.CSTTransformer): def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: if fn_name and fn_name not in original_node.name.value and fn_name != "test": return updated_node body = list(updated_node.body.body) if any(getattr(stmt, "code", "") == statement.code for stmt in body): return updated_node if operation == "insert_before": body = [statement, *body] else: body = [*body, statement] return updated_node.with_changes(body=updated_node.body.with_changes(body=body)) return module.visit(InsertTransformer()).code def _apply_textual_operation_idempotent(source: str, operation: str, identifier: str, rendered: str) -> str: if rendered in source: return source if operation in {"insert_before", "insert_after"} and identifier: match = re.search(re.escape(identifier), source) if not match: raise ValueError(f"Target identifier not found: {identifier}") insert = rendered + "\n" if operation == "insert_before": return source[: match.start()] + insert + source[match.start() :] return source[: match.end()] + "\n" + rendered + source[match.end() :] raise ValueError(f"Unsupported textual operation: {operation}") def _cst_name(node: cst.CSTNode) -> str: if isinstance(node, cst.Name): return node.value if isinstance(node, cst.Attribute): left = _cst_name(node.value) return f"{left}.{node.attr.value}" if left else node.attr.value return "" def _apply_reset_fixture(source: str, scope: str) -> str: fixture_name = "_flakeforge_reset_state" if f"def {fixture_name}(" in source: return source module = cst.parse_module(source) fixture_code = ( f"@pytest.fixture(autouse=True, scope='{scope}')\n" f"def {fixture_name}():\n" " for _key, _value in list(globals().items()):\n" " if _key.startswith('_') or _key in {'pytest', 'copy', '_flakeforge_reset_state'}:\n" " continue\n" " if isinstance(_value, dict):\n" " _value.clear()\n" " elif isinstance(_value, list):\n" " _value.clear()\n" " elif isinstance(_value, set):\n" " _value.clear()\n" " yield\n" ) try: fixture_stmt = cst.parse_statement(fixture_code) if not isinstance(fixture_stmt, cst.FunctionDef): return source except Exception: return source updated_source = source if "import pytest" not in source: updated_source = "import pytest\n" + updated_source updated_module = cst.parse_module(updated_source) body = list(updated_module.body) body.insert(0, fixture_stmt) return updated_module.with_changes(body=body).code def _apply_retry_wrapper(source: str, fn_name: str, max_attempts: int, backoff_ms: int) -> str: helper_name = "_flakeforge_retry" module = cst.parse_module(source) class RetryDecoratorAdder(cst.CSTTransformer): def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: if fn_name and fn_name not in original_node.name.value and fn_name != "test": return updated_node decorator_expr = cst.parse_expression( f"{helper_name}(max_attempts={max_attempts}, backoff_ms={backoff_ms})" ) existing = [d.decorator.code for d in updated_node.decorators] if any(helper_name in value for value in existing): return updated_node return updated_node.with_changes(decorators=[*updated_node.decorators, cst.Decorator(decorator_expr)]) transformed = module.visit(RetryDecoratorAdder()).code if f"def {helper_name}(" in transformed: return transformed helper_code = ( "\n" "def _flakeforge_retry(max_attempts: int = 2, backoff_ms: int = 100):\n" " def _decorator(fn):\n" " def _wrapped(*args, **kwargs):\n" " _last_exc = None\n" " for _attempt in range(max_attempts):\n" " try:\n" " return fn(*args, **kwargs)\n" " except Exception as _exc:\n" " _last_exc = _exc\n" " time.sleep(backoff_ms / 1000.0)\n" " raise _last_exc\n" " return _wrapped\n" " return _decorator\n" ) if "import time" not in transformed: transformed = "import time\n" + transformed return transformed + helper_code def _apply_seed_call(source: str, fn_name: str, library: str) -> str: module = cst.parse_module(source) seed_statements: List[cst.BaseStatement] = [] if library in {"random", "both"}: seed_statements.append(cst.parse_statement("random.seed(42)\n")) if library in {"numpy", "both"}: seed_statements.append(cst.parse_statement("numpy.random.seed(42)\n")) class SeedInserter(cst.CSTTransformer): def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: if fn_name and fn_name not in original_node.name.value and fn_name != "test": return updated_node body = list(updated_node.body.body) existing_codes = [stmt.code for stmt in body] to_insert: List[cst.BaseStatement] = [] for stmt in seed_statements: if stmt.code not in existing_codes: to_insert.append(stmt) if not to_insert: return updated_node return updated_node.with_changes(body=updated_node.body.with_changes(body=[*to_insert, *body])) transformed = module.visit(SeedInserter()).code if library in {"random", "both"} and "import random" not in transformed: transformed = "import random\n" + transformed if library in {"numpy", "both"} and "import numpy" not in transformed: transformed = "import numpy\n" + transformed return transformed # ── V2 Deep-Action Implementations (Bug 2 fix) ──────────────────────────────────── def _apply_refactor_concurrency(source: str, target: Dict[str, Any]) -> str: """ Swap a threading/sync primitive for a safer alternative. Production fixes: 1. Matches the FULL dotted name (e.g. "threading.Lock") — not a loose substring — so "asyncio.Lock" and "FileLock" are never accidentally swapped. 2. Also rewrites the import statement to match the new location. If `from threading import Lock` exists, changes it to the new module. If `import threading` exists and the new primitive is in asyncio, adds `import asyncio` if not already present. """ from_primitive: str = str(target.get("from_primitive", "threading.Lock")) to_primitive: str = str(target.get("to_primitive", "threading.RLock")) # Derive module + class name for import rewriting from_module = from_primitive.rsplit(".", 1)[0] if "." in from_primitive else "" to_module = to_primitive.rsplit(".", 1)[0] if "." in to_primitive else "" to_class = to_primitive.rsplit(".", 1)[-1] module = cst.parse_module(source) class PrimitiveSwapper(cst.CSTTransformer): """Rewrite every call site that exactly matches from_primitive.""" def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression: # Use EXACT full dotted name comparison, not substring. full_name = _cst_name(updated_node.func) if full_name != from_primitive: return updated_node try: replacement_func = cst.parse_expression(to_primitive) return updated_node.with_changes(func=replacement_func) except Exception: return updated_node class ImportRewriter(cst.CSTTransformer): """ Rewrite import statements to match the new primitive's module. Handles all four patterns: import threading → also imports to_module from threading import Lock → from to_module import to_class from threading import Lock as L → from to_module import to_class as L import threading as thr → also imports to_module as its alias """ def leave_ImportFrom( self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom ) -> cst.ImportFrom: if from_module and to_module and from_module != to_module: # Get the module name as a string safely try: mod_str = cst.parse_module("").code_for_node(updated_node.module) if updated_node.module else "" except Exception: mod_str = "" if mod_str == from_module: new_module = cst.parse_expression(to_module) # Rewrite the imported name new_names: List[cst.ImportAlias] = [] for alias_node in (updated_node.names if isinstance(updated_node.names, (list, tuple)) else []): if isinstance(alias_node, cst.ImportAlias): name_str = alias_node.name.value if isinstance(alias_node.name, cst.Name) else "" from_class = from_primitive.rsplit(".", 1)[-1] if name_str == from_class: new_alias = alias_node.with_changes( name=cst.Name(to_class) ) new_names.append(new_alias) else: new_names.append(alias_node) else: pass # skip star imports return updated_node.with_changes( module=new_module, names=new_names if new_names else updated_node.names, ) return updated_node # Step 1: swap all call sites transformed = module.visit(PrimitiveSwapper()).code # Step 2: rewrite imports transformed_module = cst.parse_module(transformed) transformed = transformed_module.visit(ImportRewriter()).code # Step 3: ensure the new primitive's module is imported if it isn't already if to_module and to_module not in transformed.split(".")[0]: if f"import {to_module}" not in transformed and f"from {to_module}" not in transformed: transformed = f"import {to_module}\n" + transformed return transformed def _apply_isolate_boundary(source: str, target: Dict[str, Any]) -> str: """Wrap external call in circuit-breaker / timeout.""" boundary_call = str(target.get("boundary_call", "")) pattern = str(target.get("pattern", "timeout")) module = cst.parse_module(source) class BoundaryIsolator(cst.CSTTransformer): def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.BaseExpression: name = _cst_name(updated_node.func) if boundary_call and boundary_call in (name or ""): try: if pattern == "timeout": wrapper = cst.parse_expression(f"asyncio.wait_for({name}(), timeout=5.0)") elif pattern == "circuit_breaker": wrapper = cst.parse_expression(f"_circuit_breaker({name})") else: return updated_node return wrapper except Exception: pass return updated_node return module.visit(BoundaryIsolator()).code def _apply_extract_async_scope(source: str, target: Dict[str, Any]) -> str: """ Structurally fix the async/sync boundary problem. direction='sync_to_async' (most common case): Converts a synchronous function to async AND wraps every detected blocking I/O call (DB / HTTP / socket) with ``asyncio.to_thread()`` so the event loop is never stalled. This is the correct fix for: ``async def f(): db.commit()`` direction='async_to_sync': Removes the async/await keywords so the function can be called safely from synchronous contexts. Previous implementation only flipped the `async def` keyword but left blocking calls inside untouched — which made things WORSE. """ from typing import Set as _Set fn_name = str(target.get("function_name", "")) direction = str(target.get("direction", "sync_to_async")) # Blocking signatures reused from causal_graph detection sets BLOCKING_CALL_SIGNATURES: _Set[str] = { # DB "session.commit", "session.execute", "session.add", "session.flush", "cursor.execute", "cursor.executemany", "db.commit", "db.execute", "db.session.commit", "engine.connect", "engine.execute", "collection.find", "collection.insert_one", "collection.update_one", "redis.set", "redis.get", # HTTP "requests.get", "requests.post", "requests.put", "requests.delete", "requests.patch", "urllib.request.urlopen", # Filesystem / socket blocking "socket.recv", "socket.accept", # time.sleep is a special case — always blocking "time.sleep", } # Partial prefixes that should be wrapped regardless of full name BLOCKING_PREFIXES = ( "requests.", "urllib.request.", "socket.", "cursor.", "session.", "db.", "engine.", "collection.", "redis.", ) module = cst.parse_module(source) class AsyncScopeExtractor(cst.CSTTransformer): def _is_target_function(self, node_name: str) -> bool: return (not fn_name) or (fn_name in node_name) def _call_is_blocking(self, call_node: cst.Call) -> bool: full = _cst_name(call_node.func) return ( full in BLOCKING_CALL_SIGNATURES or any(full.startswith(pfx) for pfx in BLOCKING_PREFIXES) ) # ── sync_to_async: make function async + wrap blocking calls ────────── def leave_FunctionDef( self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef ) -> cst.FunctionDef: if direction != "sync_to_async": return updated_node if not self._is_target_function(original_node.name.value): return updated_node if original_node.asynchronous: # Already async — just wrap blocking calls inside return updated_node return updated_node.with_changes( asynchronous=cst.Asynchronous(), ) def leave_Call( self, original_node: cst.Call, updated_node: cst.Call ) -> cst.BaseExpression: """Wrap blocking calls inside the target function with asyncio.to_thread().""" if direction != "sync_to_async": return updated_node if not self._call_is_blocking(updated_node): return updated_node # Build: await asyncio.to_thread(original_call) # The `await` keyword is added by the parent Await node; # here we just wrap the call in asyncio.to_thread(...). try: wrapped = cst.parse_expression( f"asyncio.to_thread({_cst_name(updated_node.func)})" ) # Preserve original arguments inside to_thread if isinstance(wrapped, cst.Call) and updated_node.args: all_args = list(wrapped.args) + list(updated_node.args) wrapped = wrapped.with_changes(args=all_args) return wrapped except Exception: return updated_node # ── async_to_sync: strip async/await ────────────────────────────────── def leave_AsyncFunctionDef( self, original_node: cst.AsyncFunctionDef, updated_node: cst.AsyncFunctionDef ) -> cst.FunctionDef: if direction != "async_to_sync": return updated_node # type: ignore[return-value] if not self._is_target_function(original_node.name.value): return updated_node # type: ignore[return-value] # Convert AsyncFunctionDef → FunctionDef sync_fn = cst.FunctionDef( name=updated_node.name, params=updated_node.params, body=updated_node.body, decorators=updated_node.decorators, returns=updated_node.returns, leading_lines=updated_node.leading_lines, lines_after_decorators=updated_node.lines_after_decorators, ) return sync_fn transformed = module.visit(AsyncScopeExtractor()).code # Ensure asyncio is imported when doing sync_to_async if direction == "sync_to_async" and "import asyncio" not in transformed: transformed = "import asyncio\n" + transformed return transformed def _apply_harden_idempotency(source: str, target: Dict[str, Any]) -> str: """Add idempotency guard to a state-mutating function.""" state_target = str(target.get("state_target", "")) key_strategy = str(target.get("key_strategy", "uuid")) module = cst.parse_module(source) idempotency_guard = ( "\n" "def _ensure_idempotency(key_strategy='uuid'):\n" " import uuid, hashlib\n" " if key_strategy == 'uuid':\n" " return str(uuid.uuid4())\n" " return hashlib.md5(str(key_strategy).encode()).hexdigest()\n" ) class IdempotencyHardener(cst.CSTTransformer): def leave_FunctionDef(self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef) -> cst.FunctionDef: if state_target and state_target not in original_node.name.value: return updated_node body = list(updated_node.body.body) guard_stmt = cst.parse_statement("_idempotency_key = _ensure_idempotency()\n") if any(getattr(s, "code", "").find("_idempotency_key") >= 0 for s in body): return updated_node return updated_node.with_changes(body=updated_node.body.with_changes(body=[guard_stmt, *body])) transformed = module.visit(IdempotencyHardener()).code if "def _ensure_idempotency" not in transformed: transformed = transformed + idempotency_guard return transformed def _parse_js_ts_summary(path: Path) -> ASTSummary: source = path.read_text(encoding="utf-8", errors="ignore") try: from tree_sitter import Language, Parser import tree_sitter_javascript as ts_js except Exception: return ASTSummary([], [], [], [], [], []) parser = Parser() language = Language(ts_js.language()) parser.language = language tree = parser.parse(bytes(source, "utf-8")) functions: List[Dict[str, Any]] = [] classes: List[Dict[str, Any]] = [] imports: List[str] = [] external_calls: List[str] = [] stack = [tree.root_node] while stack: node = stack.pop() stack.extend(node.children) if node.type in {"function_declaration", "method_definition", "arrow_function"}: name = "" for child in node.children: if child.type in {"identifier", "property_identifier"}: name = source[child.start_byte : child.end_byte] break functions.append( { "name": name, "is_async": "async" in source[node.start_byte : node.end_byte].split("(", 1)[0], "args": [], "decorators": [], "line_range": [node.start_point[0] + 1, node.end_point[0] + 1], } ) if node.type == "class_declaration": class_name = "" for child in node.children: if child.type == "identifier": class_name = source[child.start_byte : child.end_byte] break classes.append( { "name": class_name, "bases": [], "decorators": [], "line_range": [node.start_point[0] + 1, node.end_point[0] + 1], } ) if node.type == "import_statement": imports.append(source[node.start_byte : node.end_byte].strip()) if node.type == "call_expression": snippet = source[node.start_byte : node.end_byte] if any(mod in snippet for mod in ["axios", "fetch", "redis", "pg."]): external_calls.append(snippet[:120]) return ASTSummary( functions=functions, classes=classes, imports=sorted(set(imports)), global_vars=[], threading_primitives=[], external_calls=sorted(set(external_calls)), )