diff --git a/libs/openant-core/core/analyzer.py b/libs/openant-core/core/analyzer.py index f8255f13..3010532a 100644 --- a/libs/openant-core/core/analyzer.py +++ b/libs/openant-core/core/analyzer.py @@ -47,6 +47,24 @@ load_context = None +def _unit_security_classification(unit): + """Return a unit's security_classification, regardless of enhance mode. + + Agentic enhance writes ``unit['agent_context']['security_classification']``; + single-shot enhance writes ``unit['llm_context']`` and historically had no + classification at all. Reading only ``agent_context`` silently dropped every + single-shot unit from ``--exploitable-only/all``. Prefer + agent_context, fall back to llm_context, else None. + """ + agent_ctx = unit.get("agent_context") + if isinstance(agent_ctx, dict) and agent_ctx.get("security_classification") is not None: + return agent_ctx.get("security_classification") + llm_ctx = unit.get("llm_context") + if isinstance(llm_ctx, dict) and llm_ctx.get("security_classification") is not None: + return llm_ctx.get("security_classification") + return None + + def _process_unit(client, unit, index, json_corrector, app_context): """Process a single unit for Stage 1 detection. @@ -349,10 +367,23 @@ def run_analysis( keep = ("exploitable",) else: # "all" — default when filtering is enabled keep = ("exploitable", "vulnerable_internal") - units = [ - u for u in units - if u.get("agent_context", {}).get("security_classification") in keep - ] + # Read classification mode-agnostically (agent_context OR llm_context) + # so single-shot-enhanced datasets are not silently dropped. + classified = sum(1 for u in units if _unit_security_classification(u) is not None) + units = [u for u in units if _unit_security_classification(u) in keep] + # Loud guard: a filter that matches nothing because the dataset carries + # NO classifications at all is almost always an un-enhanced / wrong-mode + # input, not a genuine "0 exploitable units" result. Warn instead of + # silently returning an empty analysis. + if original_count and classified == 0: + print( + f"[Analyze] WARNING: --exploitable filter ({exploitable_filter}) " + f"requested but NONE of the {original_count} units carry a " + "security_classification (run `enhance` first; single-shot mode " + "must populate llm_context.security_classification). " + "Filter matched 0 units.", + file=sys.stderr, + ) print(f"[Analyze] Exploitable filter ({exploitable_filter}): {original_count} -> {len(units)} units", file=sys.stderr) if limit: diff --git a/libs/openant-core/core/enhancer.py b/libs/openant-core/core/enhancer.py index 70879b81..43a2d4a8 100644 --- a/libs/openant-core/core/enhancer.py +++ b/libs/openant-core/core/enhancer.py @@ -39,7 +39,7 @@ def enhance_dataset( analyzer_output_path: Path to analyzer_output.json (required for agentic mode). repo_path: Path to the repository (required for agentic mode). mode: "agentic" (thorough, tool-use) or "single-shot" (fast, cheaper). - checkpoint_path: Path to save/resume checkpoint (agentic mode only). + checkpoint_path: Path to save/resume checkpoint (both modes). If None, auto-derived from output_path. model: "sonnet" (default, cost-effective). workers: Number of parallel workers (default: 8). @@ -55,8 +55,11 @@ def enhance_dataset( print(f"[Enhance] Mode: {mode}", file=sys.stderr) print(f"[Enhance] Model: {model_id}", file=sys.stderr) - # Auto-derive checkpoint path for agentic mode - if mode == "agentic" and checkpoint_path is None: + # Auto-derive checkpoint path for BOTH modes so single-shot also resumes + # after an interrupt / cost-cap instead of reprocessing every unit + # Single-shot is the cheap, high-volume mode, so wasted + # re-enhancement there is the most costly to lose. + if checkpoint_path is None: output_dir = os.path.dirname(os.path.abspath(output_path)) checkpoint_path = os.path.join(output_dir, "enhance_checkpoints") @@ -106,6 +109,7 @@ def _on_restored(count: int): dataset, progress_callback=_on_unit_done, workers=workers, + checkpoint_path=checkpoint_path, ) else: raise ValueError(f"Unknown enhancement mode: {mode}. Use 'agentic' or 'single-shot'.") diff --git a/libs/openant-core/tests/test_enhance_resilience.py b/libs/openant-core/tests/test_enhance_resilience.py new file mode 100644 index 00000000..ec7ccec1 --- /dev/null +++ b/libs/openant-core/tests/test_enhance_resilience.py @@ -0,0 +1,203 @@ +"""Regression tests for enhance/analyzer resilience fixes. + +Covers four confirmed bugs: + + - is_retryable_error() string branch omits "529" and "overloaded" -> + Anthropic-overloaded errors never retried on the analyzer detection path + (which feeds str(e), not a dict). + - agentic post-loop transient-error retry was a single pass; a unit that + fails again on the one retry is never re-attempted. + - single-shot enhance_dataset had no checkpoint/resume, so a --checkpoint path + was silently dropped and an interrupted run reprocessed every unit. + - single-shot enhance writes only llm_context + (no agent_context.security_classification), so --exploitable-only/all + silently dropped every single-shot unit with no operator signal. + +No external network: ContextEnhancer.enhance_unit is monkeypatched. +""" +import os + +import pytest + +from utilities.rate_limiter import is_retryable_error + + +@pytest.fixture(autouse=True) +def _anthropic_api_key(monkeypatch): + """Provide a dummy API key so ContextEnhancer's default AnthropicClient + constructs without a real credential. No network call is made: every test + monkeypatches the actual enhancement methods, so the SDK client is never + used to issue a request.""" + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-key") + + +# -------------------------------------------------------------------------- +# 529 / overloaded must be retryable via the string branch. +# -------------------------------------------------------------------------- +class TestIsRetryable529: + def test_529_overloaded_string_is_retryable(self): + # The exact shape the Anthropic SDK stringifies a 529 into. The + # analyzer detection path stores str(e), forcing the string branch. + err = "Error code: 529 - {'type': 'error', 'error': {'type': 'overloaded_error', 'message': 'Overloaded'}}" + assert is_retryable_error(err) is True + + def test_overloaded_word_alone_is_retryable(self): + assert is_retryable_error("server overloaded, try again") is True + + def test_existing_5xx_strings_still_retryable(self): + for code in ("500", "502", "503", "504"): + assert is_retryable_error(f"Error code: {code} - server error") is True + + def test_client_error_400_still_not_retryable(self): + # Guard against over-broadening: a 400 must remain non-retryable. + assert is_retryable_error("Error code: 400 - bad request") is False + + def test_structured_529_dict_still_retryable(self): + assert is_retryable_error({"type": "api_status", "status_code": 529}) is True + + +# -------------------------------------------------------------------------- +# Shared helpers for the ContextEnhancer-level tests. +# -------------------------------------------------------------------------- +def _make_enhancer(): + from utilities.context_enhancer import ContextEnhancer + + return ContextEnhancer(client=None, tracker=None) + + +def _dataset(n=3): + return {"units": [{"id": f"u{i}", "code": f"def f{i}(): pass"} for i in range(n)]} + + +# -------------------------------------------------------------------------- +# agentic retry must be a bounded multi-round loop. +# -------------------------------------------------------------------------- +class TestAgenticBoundedRetry: + def test_unit_failing_on_first_retry_is_retried_again(self, monkeypatch, tmp_path): + """A transient error that persists for one retry must get another round + (up to the cap), instead of being abandoned after a single pass.""" + from unittest.mock import MagicMock + + from utilities import context_enhancer as ce + + enh = _make_enhancer() + # Avoid building a real repository index (no analyzer output / repo). + fake_index = MagicMock() + fake_index.get_statistics.return_value = { + "total_functions": 0, "total_files": 0, + } + monkeypatch.setattr(ce, "load_index_from_file", lambda *a, **k: fake_index) + + # attempts[uid] = how many times we've enhanced that unit. + attempts = {} + + def fake_agent(unit, index, tracker, verbose, client=None): + uid = unit.get("id") + attempts[uid] = attempts.get(uid, 0) + 1 + # u1 stays transiently-broken for the first 2 attempts, then succeeds. + if uid == "u1" and attempts[uid] < 3: + unit["agent_context"] = { + "error": {"type": "api_status", "status_code": 529}, + "security_classification": "neutral", + "confidence": 0.0, + } + else: + unit["agent_context"] = { + "security_classification": "neutral", + "confidence": 0.5, + } + + monkeypatch.setattr(ce, "enhance_unit_with_agent", fake_agent) + + ds = _dataset(2) + result = enh.enhance_dataset_agentic( + dataset=ds, + analyzer_output_path=None, + repo_path=None, + workers=1, + ) + + # With a bounded multi-round retry, u1 is attempted 3 times total + # (initial + 2 retry rounds) and ends without an error. + u1 = next(u for u in result["units"] if u["id"] == "u1") + assert attempts["u1"] >= 3, f"u1 only attempted {attempts['u1']}x (single-pass retry bug)" + assert not u1.get("agent_context", {}).get("error"), "u1 still errored after bounded retry" + + +# -------------------------------------------------------------------------- +# single-shot enhance must support checkpoint/resume. +# -------------------------------------------------------------------------- +class TestSingleShotCheckpoint: + def test_enhance_dataset_accepts_checkpoint_path(self, monkeypatch, tmp_path): + """Single-shot enhance_dataset must accept a checkpoint dir and persist + per-unit results so an interrupted run can resume.""" + from utilities.context_enhancer import ContextEnhancer + + enh = ContextEnhancer(client=None, tracker=None) + + def fake_enhance_unit(unit, units_by_id): + unit["llm_context"] = { + "reasoning": "ok", + "confidence": 0.9, + "security_classification": "neutral", + } + return unit + + monkeypatch.setattr(enh, "enhance_unit", fake_enhance_unit) + + cp_dir = str(tmp_path / "enhance_checkpoints") + enh.enhance_dataset(_dataset(3), workers=1, checkpoint_path=cp_dir) + + files = [f for f in os.listdir(cp_dir) if f.endswith(".json")] + assert len(files) == 3, f"expected 3 per-unit checkpoints, found {files}" + + def test_resume_skips_already_completed_units(self, monkeypatch, tmp_path): + """On resume, units already checkpointed must not be re-enhanced.""" + from utilities.context_enhancer import ContextEnhancer + + cp_dir = str(tmp_path / "enhance_checkpoints") + + # First run: complete all units. + enh1 = ContextEnhancer(client=None, tracker=None) + monkeypatch.setattr( + enh1, "enhance_unit", + lambda unit, by_id: unit.update( + {"llm_context": {"reasoning": "ok", "confidence": 0.9}} + ) or unit, + ) + enh1.enhance_dataset(_dataset(3), workers=1, checkpoint_path=cp_dir) + + # Second run: count how many units get (re-)enhanced. + enh2 = ContextEnhancer(client=None, tracker=None) + seen = [] + + def counting_enhance(unit, by_id): + seen.append(unit.get("id")) + unit["llm_context"] = {"reasoning": "ok", "confidence": 0.9} + return unit + + monkeypatch.setattr(enh2, "enhance_unit", counting_enhance) + enh2.enhance_dataset(_dataset(3), workers=1, checkpoint_path=cp_dir) + + assert seen == [], f"resume re-enhanced units that were already done: {seen}" + + +# -------------------------------------------------------------------------- +# exploitable filter must not silently drop single-shot units. +# -------------------------------------------------------------------------- +class TestExploitableFilterSingleShot: + def test_filter_reads_llm_context_classification_fallback(self): + """The analyzer exploitable filter must fall back to llm_context's + security_classification when agent_context is absent (single-shot).""" + from core.analyzer import _unit_security_classification + + # single-shot unit: classification lives under llm_context + u = {"id": "x", "llm_context": {"security_classification": "exploitable"}} + assert _unit_security_classification(u) == "exploitable" + + # agentic unit: classification under agent_context still works + a = {"id": "y", "agent_context": {"security_classification": "vulnerable_internal"}} + assert _unit_security_classification(a) == "vulnerable_internal" + + # no classification anywhere + assert _unit_security_classification({"id": "z"}) is None diff --git a/libs/openant-core/utilities/context_enhancer.py b/libs/openant-core/utilities/context_enhancer.py index 2f7dea20..e0f9f165 100644 --- a/libs/openant-core/utilities/context_enhancer.py +++ b/libs/openant-core/utilities/context_enhancer.py @@ -48,6 +48,11 @@ def _get_step_checkpoint(): # Use Sonnet for context enhancement (cost-effective auxiliary task) CONTEXT_ENHANCEMENT_MODEL = "claude-sonnet-4-20250514" +# Max bounded rounds of the post-loop transient-error retry. Each round +# re-attempts units still carrying a retryable error; rounds stop early once +# none remain or a round recovers nothing. +MAX_RETRY_ROUNDS = 3 + def _build_error_info(exc: Exception) -> dict: """Build a structured error dict from an exception. @@ -326,18 +331,26 @@ def enhance_dataset( batch_size: int = 10, progress_callback: Optional[Callable] = None, workers: int = 10, + checkpoint_path: str = None, ) -> dict: """ Enhance all units in a dataset (single-shot mode). Uses ThreadPoolExecutor for parallel processing when workers > 1. + Supports checkpoint/resume symmetrically with the agentic path + when ``checkpoint_path`` is provided, each completed + unit's ``llm_context`` is saved to its own file under that directory and + an interrupted run resumes without re-enhancing finished units. + Args: dataset: The dataset from unit_generator.js batch_size: Number of units to process before printing progress progress_callback: Optional callback(unit_id, classification, unit_elapsed) called after each unit completes. workers: Number of parallel workers (default: 10). + checkpoint_path: Path to a checkpoint directory (enables resume). + When None, no checkpoints are written (prior behavior). Returns: Enhanced dataset @@ -354,13 +367,41 @@ def enhance_dataset( total = len(units) + # Checkpoint directory setup + resume (mirror enhance_dataset_agentic). + checkpoint_dir = None + processed_ids = set() + if checkpoint_path: + checkpoint_dir = ( + checkpoint_path + if os.path.isdir(checkpoint_path) or not checkpoint_path.endswith(".json") + else os.path.splitext(checkpoint_path)[0] + "_checkpoints" + ) + os.makedirs(checkpoint_dir, exist_ok=True) + processed_ids = self._load_completed_units(checkpoint_dir, context_key="llm_context") + # Restore llm_context for already-completed units. + for unit in units: + unit_id = unit.get("id") + if unit_id in processed_ids: + cp_file = os.path.join(checkpoint_dir, f"{self._safe_filename(unit_id)}.json") + if os.path.exists(cp_file): + cp_data = read_json(cp_file) + unit["llm_context"] = cp_data.get("llm_context", {}) + if "code" in cp_data: + unit["code"] = cp_data["code"] + if processed_ids: + self._log("info", + f"Restored {len(processed_ids)} already-processed units from checkpoints", + units=len(processed_ids)) + self._log("info", f"Enhancing {total} units with LLM context (single-shot mode)", units=total) self._log("info", f"Model: {CONTEXT_ENHANCEMENT_MODEL}") mode = "sequential" if workers <= 1 else f"parallel ({workers} workers)" self._log("info", f"Mode: {mode}") - # Build lookup dict for context gathering + # Build lookup dict for context gathering (over ALL units, so cross-unit + # context lookup still works even when some are restored/skipped). units_by_id = {u.get("id"): u for u in units} + units_to_process = [u for u in units if u.get("id") not in processed_ids] def _process_one(unit): """Process a single unit. Mutates unit in-place.""" @@ -370,20 +411,33 @@ def _process_one(unit): ctx = unit.get("llm_context", {}) classification = ctx.get("confidence", "unknown") worker = threading.current_thread().name + if checkpoint_dir: + self._save_unit_checkpoint(unit, checkpoint_dir, context_key="llm_context") return unit.get("id", "?"), str(classification), unit_elapsed, worker if workers <= 1: - for unit in units: - uid, classification, elapsed, _worker = _process_one(unit) - if progress_callback: - progress_callback(uid, classification, elapsed) + try: + for unit in units_to_process: + uid, classification, elapsed, _worker = _process_one(unit) + if progress_callback: + progress_callback(uid, classification, elapsed) + except KeyboardInterrupt: + self._log("warning", "Interrupted — progress saved to checkpoints") + return dataset else: - with ThreadPoolExecutor(max_workers=workers) as executor: - futures = {executor.submit(_process_one, unit): unit for unit in units} + executor = ThreadPoolExecutor(max_workers=workers) + futures = {executor.submit(_process_one, unit): unit for unit in units_to_process} + try: for future in as_completed(futures): uid, classification, elapsed, worker = future.result() if progress_callback: progress_callback(uid, f"{classification} [{worker}]", elapsed) + except KeyboardInterrupt: + self._log("warning", "Interrupted — cancelling pending work...") + executor.shutdown(wait=False, cancel_futures=True) + self._log("info", "Progress saved to checkpoints") + return dataset + executor.shutdown(wait=False) # Recompute stats from unit results (thread-safe) self.stats = { @@ -675,23 +729,35 @@ def _update_summary(classification, unit): return dataset executor.shutdown(wait=False) - # Auto-retry failed units with transient errors (rate limit, connection, timeout, 5xx) - retryable_units = [ - (i, unit) for i, unit in enumerate(units) - if is_retryable_error(unit.get("agent_context", {}).get("error")) - ] - if retryable_units: + # Auto-retry failed units with transient errors (rate limit, connection, + # timeout, 5xx). This is a BOUNDED MULTI-ROUND loop: a unit that fails + # again with a still-transient error gets re-attempted on the next round + # (up to MAX_RETRY_ROUNDS). The previous single-pass version abandoned a + # unit after one retry even if its error was still transient + # Rounds stop early once no retryable + # units remain or no progress is made. + for _retry_round in range(MAX_RETRY_ROUNDS): + retryable_units = [ + (i, unit) for i, unit in enumerate(units) + if is_retryable_error(unit.get("agent_context", {}).get("error")) + ] + if not retryable_units: + break + rate_limiter = get_rate_limiter() backoff = rate_limiter.time_until_ready() if backoff > 0: self._log("info", - f"Retrying {len(retryable_units)} failed units " + f"Retry round {_retry_round + 1}/{MAX_RETRY_ROUNDS}: " + f"{len(retryable_units)} failed units " f"(waiting {backoff:.0f}s for rate limit to clear)...") rate_limiter.wait_if_needed() else: self._log("info", - f"Retrying {len(retryable_units)} failed units (transient errors)...") + f"Retry round {_retry_round + 1}/{MAX_RETRY_ROUNDS}: " + f"{len(retryable_units)} failed units (transient errors)...") + round_recovered = 0 # Retry sequentially to avoid re-triggering rate limit for i, unit in retryable_units: # Clear previous error @@ -700,6 +766,7 @@ def _update_summary(classification, unit): # Update summary: retry succeeded → flip error to completed if classification != "error": + round_recovered += 1 _summary_errors = max(0, _summary_errors - 1) _summary_completed += 1 # Decrement the old error type count (best effort) @@ -723,6 +790,11 @@ def _update_summary(classification, unit): if progress_callback: progress_callback(uid, f"{classification} (retry)", elapsed) + # No unit recovered this round → further rounds would just repeat the + # same persistent failures. Stop to avoid burning quota. + if round_recovered == 0: + break + # Write final summary with phase="done" if _summary_cp is not None: _summary_cp.write_summary(total, _summary_completed, _summary_errors, @@ -771,20 +843,27 @@ def _safe_filename(unit_id: str) -> str: from utilities.safe_filename import safe_filename return safe_filename(unit_id) - def _save_unit_checkpoint(self, unit: dict, checkpoint_dir: str): - """Save a single unit's result to its own checkpoint file.""" + def _save_unit_checkpoint(self, unit: dict, checkpoint_dir: str, context_key: str = "agent_context"): + """Save a single unit's result to its own checkpoint file. + + ``context_key`` selects which enhancement context to persist: + ``agent_context`` for agentic mode, ``llm_context`` for single-shot. + The key is recorded so resume knows where to restore. + """ unit_id = unit.get("id", "unknown") filename = self._safe_filename(unit_id) + ".json" filepath = os.path.join(checkpoint_dir, filename) + ctx = unit.get(context_key, {}) cp_data = { "id": unit_id, - "agent_context": unit.get("agent_context", {}), + "context_key": context_key, + context_key: ctx, } # Include code if it was modified by the agent if "code" in unit: cp_data["code"] = unit["code"] - # Include per-unit usage from agent_metadata - meta = cp_data["agent_context"].get("agent_metadata", {}) + # Include per-unit usage from agent_metadata (agentic only) + meta = ctx.get("agent_metadata", {}) if isinstance(ctx, dict) else {} if meta.get("input_tokens") or meta.get("output_tokens"): cp_data["usage"] = { "input_tokens": meta.get("input_tokens", 0), @@ -793,8 +872,13 @@ def _save_unit_checkpoint(self, unit: dict, checkpoint_dir: str): } write_json(filepath, cp_data) - def _load_completed_units(self, checkpoint_dir: str) -> set: - """Load the set of completed unit IDs from per-unit checkpoint files.""" + def _load_completed_units(self, checkpoint_dir: str, context_key: str = "agent_context") -> set: + """Load the set of completed unit IDs from per-unit checkpoint files. + + A unit counts as completed if its persisted context (for the relevant + ``context_key``) is present and carries no error. Older agentic-only + checkpoints have no ``context_key`` field and are read as agent_context. + """ completed = set() if not os.path.isdir(checkpoint_dir): return completed @@ -805,8 +889,9 @@ def _load_completed_units(self, checkpoint_dir: str) -> set: try: cp_data = read_json(filepath) unit_id = cp_data.get("id") - agent_ctx = cp_data.get("agent_context", {}) - if unit_id and agent_ctx and not agent_ctx.get("error"): + key = cp_data.get("context_key", context_key) + ctx = cp_data.get(key, {}) + if unit_id and ctx and not (isinstance(ctx, dict) and ctx.get("error")): completed.add(unit_id) except (json.JSONDecodeError, OSError): continue diff --git a/libs/openant-core/utilities/rate_limiter.py b/libs/openant-core/utilities/rate_limiter.py index 3416f1b1..aaf03993 100644 --- a/libs/openant-core/utilities/rate_limiter.py +++ b/libs/openant-core/utilities/rate_limiter.py @@ -236,8 +236,17 @@ def is_retryable_error(error_info: dict | str | None) -> bool: return False - # String-based error checking + # String-based error checking. + # NOTE: the analyzer detection path (core/analyzer.py) stores the raw + # str(e) of an exception rather than a structured dict, so an Anthropic + # HTTP 529 surfaces here as e.g. + # "Error code: 529 - {...'type':'overloaded_error'...}" + # 529 ("overloaded") is the most common transient Anthropic failure under + # load and is a 5xx, so it must be retried. The structured dict branch + # above already retries it via status_code >= 500; mirror that here so the + # string path is not silently non-retryable. error_str = str(error_info).lower() return any(term in error_str for term in ( - "rate_limit", "connection", "timeout", "500", "502", "503", "504" + "rate_limit", "connection", "timeout", + "500", "502", "503", "504", "529", "overloaded", ))