Refactor budget clipping logic in pipeline service

- Updated `_clip_for_budget` function to return both selected test cases and remaining unselected cases for better budget management.
- Enhanced documentation for the function to clarify its purpose and return values.
- Adjusted the `run_agent_pipeline` function to accommodate the new return structure, ensuring proper handling of remaining cases.

Author: Av7danger

src/secnodeapi/services/pipeline.py

@@ -246,18 +246,28 @@ def _clip_for_budget(
     test_cases: List[TestCase],
     remaining_budget: int,
     per_endpoint_budget: int,
-) -> List[TestCase]:
+) -> Tuple[List[TestCase], List[TestCase]]:
+    """
+    Select a batch constrained by global and per-endpoint budgets.
+
+    Returns both:
+    - selected cases to execute now
+    - remaining queue preserving unselected cases for later iterations
+    """
     clipped: List[TestCase] = []
+    remaining: List[TestCase] = []
     endpoint_counts: Dict[str, int] = {}
     for case in test_cases:
         if len(clipped) >= remaining_budget:
-            break
+            remaining.append(case)
+            continue
         endpoint_count = endpoint_counts.get(case.endpoint, 0)
         if endpoint_count >= per_endpoint_budget:
+            remaining.append(case)
             continue
         endpoint_counts[case.endpoint] = endpoint_count + 1
         clipped.append(case)
-    return clipped
+    return clipped, remaining
 
 
 def _merge_unique_findings(existing: List[Finding], new_findings: List[Finding]) -> List[Finding]:
@@ -294,10 +304,11 @@ async def run_agent_pipeline(
 
     while queue and remaining_budget > 0 and iteration < pipeline_input.max_iterations:
         iteration += 1
-        batch = _clip_for_budget(queue, remaining_budget, pipeline_input.per_endpoint_budget)
+        batch, queue = _clip_for_budget(
+            queue, remaining_budget, pipeline_input.per_endpoint_budget
+        )
         if not batch:
             break
-        queue = queue[len(batch) :]
 
         results, stats = await execute_proactive_tests_detailed(
             test_cases=batch,

tests/test_pipeline_service.py

@@ -167,3 +167,104 @@ async def fake_classify(results):
     assert metrics["iterations"] == 2
     assert len(confirmed) >= 1
     assert any("CHAIN-BOLA" in test_id for batch in observed_batches[1:] for test_id in batch)
+
+
+@pytest.mark.asyncio
+async def test_run_agent_pipeline_preserves_skipped_cases_without_reexecution(monkeypatch) -> None:
+    structure = SchemaStructure(
+        title="Queue API",
+        version="1.0",
+        base_url="https://api.example.com",
+        endpoints=[APIEndpoint(path="/a", method="GET"), APIEndpoint(path="/b", method="GET")],
+        auth_schemes={},
+    )
+    seed_cases = [
+        TestCase(
+            id="A-1",
+            name="a1",
+            description="",
+            owasp_category="API9",
+            endpoint="/a",
+            method="GET",
+            params={"variant": 1},
+        ),
+        TestCase(
+            id="A-2",
+            name="a2",
+            description="",
+            owasp_category="API9",
+            endpoint="/a",
+            method="GET",
+            params={"variant": 2},
+        ),
+        TestCase(
+            id="A-3",
+            name="a3",
+            description="",
+            owasp_category="API9",
+            endpoint="/a",
+            method="GET",
+            params={"variant": 3},
+        ),
+        TestCase(
+            id="B-1",
+            name="b1",
+            description="",
+            owasp_category="API9",
+            endpoint="/b",
+            method="GET",
+        ),
+    ]
+
+    executed_ids = []
+
+    async def fake_artifacts(_):
+        return structure, seed_cases
+
+    async def fake_execute(**kwargs):
+        cases = kwargs["test_cases"]
+        executed_ids.extend(case.id for case in cases)
+        results = [
+            TestResult(
+                test_case=case,
+                status_code=200,
+                response_body='{"ok": true}',
+                response_headers={},
+                request_url=f"https://api.example.com{case.endpoint}",
+                request_headers={},
+                request_body=None,
+                response_time_ms=1.0,
+            )
+            for case in cases
+        ]
+        stats = ExecutionStats(
+            attempted=len(cases),
+            successful_requests=len(cases),
+            failed_requests=0,
+        )
+        return results, stats
+
+    async def fake_classify(_):
+        return [], []
+
+    monkeypatch.setattr("secnodeapi.services.pipeline.build_pipeline_artifacts", fake_artifacts)
+    monkeypatch.setattr("secnodeapi.services.pipeline.execute_proactive_tests_detailed", fake_execute)
+    monkeypatch.setattr("secnodeapi.services.pipeline.classify_findings", fake_classify)
+    monkeypatch.setattr("secnodeapi.services.pipeline._build_discovery_tests", lambda _: [])
+    monkeypatch.setattr("secnodeapi.services.pipeline._build_chain_tests", lambda *_: [])
+
+    pipeline_input = pipeline.PipelineInput(
+        target="https://api.example.com/swagger.json",
+        concurrency=2,
+        auth_headers={},
+        proxy=None,
+        verify_ssl=True,
+        request_budget=6,
+        per_endpoint_budget=2,
+        max_iterations=2,
+    )
+
+    _, _, _, metrics = await pipeline.run_agent_pipeline(pipeline_input)
+    assert metrics["iterations"] == 2
+    assert executed_ids.count("B-1::default") == 1
+    assert "A-3::default" in executed_ids