fix(cli): add --limit to the enhance command across all three layers

apps/openant-cli/cmd/enhance.go

@@ -32,6 +32,7 @@ var (
 	enhanceCheckpoint     string
 	enhanceWorkers        int
 	enhanceBackoff        int
+	enhanceLimit          int
 )
 
 func init() {
@@ -42,6 +43,7 @@ func init() {
 	enhanceCmd.Flags().StringVar(&enhanceCheckpoint, "checkpoint", "", "Path to save/resume checkpoint (agentic mode)")
 	enhanceCmd.Flags().IntVar(&enhanceWorkers, "workers", 8, "Number of parallel workers for LLM steps (default: 8)")
 	enhanceCmd.Flags().IntVar(&enhanceBackoff, "backoff", 30, "Seconds to wait when rate-limited (default: 30)")
+	enhanceCmd.Flags().IntVar(&enhanceLimit, "limit", 0, "Max units to enhance (0 = no limit)")
 }
 
 func runEnhance(cmd *cobra.Command, args []string) {
@@ -104,6 +106,9 @@ func runEnhance(cmd *cobra.Command, args []string) {
 	if enhanceBackoff != 30 {
 		pyArgs = append(pyArgs, "--backoff", fmt.Sprintf("%d", enhanceBackoff))
 	}
+	if enhanceLimit > 0 {
+		pyArgs = append(pyArgs, "--limit", fmt.Sprintf("%d", enhanceLimit))
+	}
 
 	result, err := python.Invoke(rt.Path, pyArgs, "", quiet, requireAPIKey())
 	if err != nil {

apps/openant-cli/cmd/enhance_limit_test.go

@@ -0,0 +1,15 @@
+package cmd
+
+import "testing"
+
+// The `enhance` command must expose a `--limit` flag for
+// parity with `analyze`/`scan`, so a user can cost-limit / test-run enhancement on N units.
+func TestEnhanceCommandHasLimitFlag(t *testing.T) {
+	f := enhanceCmd.Flags().Lookup("limit")
+	if f == nil {
+		t.Fatal("enhance command is missing the --limit flag (parity with analyze/scan)")
+	}
+	if f.DefValue != "0" {
+		t.Errorf("--limit default = %q, want \"0\" (0 = no limit)", f.DefValue)
+	}
+}

libs/openant-core/core/enhancer.py

@@ -30,6 +30,7 @@ def enhance_dataset(
     model: str = "sonnet",
     workers: int = 8,
     backoff_seconds: int = 30,
+    limit: int | None = None,
 ) -> EnhanceResult:
     """Enhance a parsed dataset with security context.
 
@@ -44,6 +45,7 @@ def enhance_dataset(
         model: "sonnet" (default, cost-effective).
         workers: Number of parallel workers (default: 8).
         backoff_seconds: Seconds to wait on rate limit before retry (default: 30).
+        limit: Max number of units to enhance (None = all). Mirrors `analyze --limit`.
 
     Returns:
         EnhanceResult with output path, stats, and usage.
@@ -72,6 +74,9 @@ def enhance_dataset(
     print(f"[Enhance] Loading dataset: {dataset_path}", file=sys.stderr)
     dataset = read_json(dataset_path)
     units = dataset.get("units", [])
+    if limit:
+        units = units[:limit]
+        dataset["units"] = units  # so the agentic/single-shot paths enhance only these
     print(f"[Enhance] Units to enhance: {len(units)}", file=sys.stderr)
 
     # Set up progress reporter

libs/openant-core/openant/cli.py

@@ -189,6 +189,7 @@ def cmd_enhance(args):
                 checkpoint_path=args.checkpoint,
                 workers=args.workers,
                 backoff_seconds=args.backoff,
+                limit=args.limit,
             )
 
             ctx.summary = {
@@ -1049,6 +1050,7 @@ def main():
     enhance_p.add_argument("--repo-path", help="Path to the repository (required for agentic mode)")
     enhance_p.add_argument("--output", "-o", help="Output path for enhanced dataset (default: {input}_enhanced.json)")
     enhance_p.add_argument("--checkpoint", help="Path to save/resume checkpoint (agentic mode)")
+    enhance_p.add_argument("--limit", type=int, help="Max units to enhance")
     enhance_p.add_argument(
         "--mode",
         choices=["agentic", "single-shot"],

libs/openant-core/tests/test_enhance_limit.py

@@ -0,0 +1,82 @@
+"""Regression tests: `enhance` lacked a `--limit` (max-units) capability at all
+three layers (Go CLI, Python CLI, core), while `analyze`/`scan` have it.
+
+These tests cover the two Python layers behaviorally:
+- the core `enhance_dataset(limit=N)` actually slices the dataset to N units (the LLM is stubbed);
+- the CLI path `enhance ... --limit N` (parser `enhance_p` + `cmd_enhance`) forwards N to enhance_dataset.
+
+(The Go layer is covered by apps/openant-cli/cmd/enhance_limit_test.go.)
+"""
+import sys
+
+import core.enhancer as enh
+import utilities.context_enhancer as _ce
+import utilities.llm_client as _llm
+
+
+_received = {}
+
+
+class _DummyEnhancer:
+    def __init__(self, **kw):
+        pass
+
+    def enhance_dataset(self, dataset, progress_callback=None, workers=8):
+        _received["units"] = len(dataset.get("units", []))
+        return dataset
+
+    def enhance_dataset_agentic(self, dataset, **kw):
+        _received["units"] = len(dataset.get("units", []))
+        return dataset
+
+
+def test_enhance_dataset_limit_slices_units(tmp_path, monkeypatch):
+    five = [{"id": f"u{i}", "llm_context": {}} for i in range(5)]
+    monkeypatch.setattr(enh, "read_json", lambda p: {"units": list(five)})
+    monkeypatch.setattr(enh, "write_json", lambda p, d: None)
+    monkeypatch.setattr(enh, "configure_rate_limiter", lambda **k: None)
+    monkeypatch.setattr(_llm, "AnthropicClient", lambda **k: object())
+    monkeypatch.setattr(_llm, "get_global_tracker", lambda: None)
+    monkeypatch.setattr(_ce, "ContextEnhancer", _DummyEnhancer)
+
+    _received.clear()
+    res = enh.enhance_dataset(
+        dataset_path=str(tmp_path / "x.json"),
+        output_path=str(tmp_path / "out.json"),
+        mode="single-shot",
+        limit=2,
+    )
+    # The enhancer must RECEIVE only the limited units -- this guards the load-bearing
+    # `dataset["units"] = units` slice (a local-only slice would leave 5 units reaching the enhancer).
+    assert _received["units"] == 2, f"enhancer received {_received.get('units')} units, expected 2"
+    assert res.units_enhanced == 2, f"limit=2 should enhance 2 of 5 units, got {res.units_enhanced}"
+
+
+def test_enhance_cli_forwards_limit_end_to_end(tmp_path, monkeypatch):
+    from openant import cli
+
+    captured = {}
+
+    def _fake_enhance_dataset(**kwargs):
+        captured.update(kwargs)
+
+        class _R:
+            enhanced_dataset_path = str(tmp_path / "out.json")
+            units_enhanced = 3
+            error_count = 0
+            classifications: dict = {}
+            error_summary: dict = {}
+
+            def to_dict(self):
+                return {"units_enhanced": 3}
+
+        return _R()
+
+    # cmd_enhance does `from core.enhancer import enhance_dataset` at call time -> patch the source attr
+    monkeypatch.setattr(enh, "enhance_dataset", _fake_enhance_dataset)
+    monkeypatch.setattr(sys, "argv", ["openant", "enhance", str(tmp_path / "x.json"), "--limit", "3"])
+
+    rc = cli.main()
+    assert rc == 0, f"enhance command returned {rc}"
+    assert captured.get("limit") == 3, \
+        f"`enhance --limit 3` must forward limit=3 to enhance_dataset, got {captured.get('limit')!r}"