From c3e8031999c973675ac2728dc202ffb6d83839e9 Mon Sep 17 00:00:00 2001 From: Evgeny Kiriyak <224408464+evkir@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:53:44 +0200 Subject: [PATCH 1/4] fix(recon): unify subdomain result shape across sync and async paths The sync ReconAgent called dns_tool.detect_subdomains, returning {subdomains: [str]}, while AsyncOrchestrator called subdomain_enum.enumerate_subdomains_async, returning {found: [{fqdn}]}. Both wrote the same recon.subdomains KB key, so any consumer of that key was correct for exactly one of the two pipeline paths. The sync agent also used the weaker of the two enumerators (11 hardcoded prefixes, sequential). Move the sync agent and the MCP subdomain_enum tool onto enumerate_subdomains and add fqdns() as the single place that derives bare hostnames for ReconResult.subdomains, so the shape cannot drift again. --- cyberai/agents/recon/agent.py | 9 ++-- cyberai/agents/recon/subdomain_enum.py | 12 +++++ cyberai/core/orchestrator.py | 5 +- cyberai/mcp/tools.py | 6 +-- .../test_plan_web_order_pipeline.py | 2 +- tests/unit/test_behavioral_probe.py | 4 +- tests/unit/test_recon.py | 46 ++++++++++++++++++- tests/unit/test_recon_guards.py | 2 +- tests/unit/test_subdomain_enum.py | 32 +++++++++++++ tests/unit/test_web_wiring.py | 4 +- 10 files changed, 104 insertions(+), 18 deletions(-) diff --git a/cyberai/agents/recon/agent.py b/cyberai/agents/recon/agent.py index 5017641..57ce001 100644 --- a/cyberai/agents/recon/agent.py +++ b/cyberai/agents/recon/agent.py @@ -8,8 +8,9 @@ from cyberai.core.scan_session import Severity from cyberai.core.types import OpenPort, ReconResult -from .dns_tool import detect_subdomains, run_dns, run_whois +from .dns_tool import run_dns, run_whois from .llm_detector import detect_llm_endpoints +from .subdomain_enum import enumerate_subdomains, fqdns from .web_surface import discover_surface from .behavioral import BehavioralFingerprint from .behavioral_probe import build_probe_context @@ -54,7 +55,7 @@ def _register_tools(self) -> None: Tool( name="subdomain_scan", description="Subdomain bruteforce", - func=detect_subdomains, + func=enumerate_subdomains, parameters={"target": "str"}, ) ) @@ -113,7 +114,7 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str # 4. Subdomains self._check_iteration_limit() - sub_result = detect_subdomains(target) + sub_result = enumerate_subdomains(target) self.kb.set("recon.subdomains", sub_result, agent=self.AGENT_NAME) results["recon.subdomains"] = sub_result self._log("subdomain_scan complete", sub_result) @@ -161,7 +162,7 @@ def run(self, target: str, context: Optional[Dict[str, Any]] = None) -> Dict[str ports=[OpenPort(**p) for p in ports if isinstance(p, dict)], whois=whois_result if isinstance(whois_result, dict) else {}, dns=dns_result if isinstance(dns_result, dict) else {}, - subdomains=(sub_result.get("subdomains", []) if isinstance(sub_result, dict) else []), + subdomains=fqdns(sub_result if isinstance(sub_result, dict) else None), ) self.kb.set("recon.result", recon_result.model_dump(), agent=self.AGENT_NAME) diff --git a/cyberai/agents/recon/subdomain_enum.py b/cyberai/agents/recon/subdomain_enum.py index 17ed27a..deafde2 100644 --- a/cyberai/agents/recon/subdomain_enum.py +++ b/cyberai/agents/recon/subdomain_enum.py @@ -64,6 +64,18 @@ ] +def fqdns(result: Dict[str, Any] | None = None) -> List[str]: + """Extract plain hostnames from an enumerate_subdomains* result. + + Both enumerators return {"found": [{"fqdn": ...}]}. Callers that need + bare hostnames (ReconResult.subdomains, the planner KB graph) must use + this instead of re-deriving the shape, or the sync and async pipeline + paths drift apart silently. + """ + found = (result or {}).get("found") or [] + return [r["fqdn"] for r in found if isinstance(r, dict) and r.get("fqdn")] + + def enumerate_subdomains( domain: str, wordlist: List[str] = None, diff --git a/cyberai/core/orchestrator.py b/cyberai/core/orchestrator.py index 858ce0b..48c24a8 100644 --- a/cyberai/core/orchestrator.py +++ b/cyberai/core/orchestrator.py @@ -401,6 +401,7 @@ async def _dispatch_async(self, session: ScanSession, phase: ScanPhase) -> Dict[ async def _run_recon_async(self, session: ScanSession) -> Dict: from cyberai.agents.recon.async_agent import AsyncReconAgent from cyberai.agents.recon.llm_detector import detect_llm_endpoints + from cyberai.agents.recon.subdomain_enum import fqdns from cyberai.core.types import OpenPort, ReconResult agent = AsyncReconAgent() @@ -423,9 +424,7 @@ async def _run_recon_async(self, session: ScanSession) -> Dict: nmap = result.get("nmap") if isinstance(result.get("nmap"), dict) else {} raw_ports = nmap.get("ports", []) if isinstance(nmap, dict) else [] subs = result.get("subdomains") if isinstance(result.get("subdomains"), dict) else {} - subdomains = [ - r["fqdn"] for r in (subs.get("found") or []) if isinstance(r, dict) and r.get("fqdn") - ] + subdomains = fqdns(subs) recon_result = ReconResult( target=session.target, ports=[OpenPort(**p) for p in raw_ports if isinstance(p, dict)], diff --git a/cyberai/mcp/tools.py b/cyberai/mcp/tools.py index 6603d44..fe66a5d 100644 --- a/cyberai/mcp/tools.py +++ b/cyberai/mcp/tools.py @@ -36,10 +36,10 @@ def register( # ── recon tools ─────────────────────────────────────────────────────── from cyberai.agents.recon.dns_tool import ( # noqa: E402 - detect_subdomains, run_dns, run_whois, ) +from cyberai.agents.recon.subdomain_enum import enumerate_subdomains # noqa: E402 from cyberai.agents.recon.nmap_tool import run_nmap # noqa: E402 register( @@ -88,8 +88,8 @@ def register( def _subdomain_handler(target: str, wordlist: list[str] | None = None) -> dict: - """Adapt detect_subdomains to MCP (wordlist optional).""" - return detect_subdomains(target, wordlist) + """Adapt enumerate_subdomains to MCP (wordlist optional).""" + return enumerate_subdomains(target, wordlist) register( diff --git a/tests/integration/test_plan_web_order_pipeline.py b/tests/integration/test_plan_web_order_pipeline.py index fabef85..9d531e3 100644 --- a/tests/integration/test_plan_web_order_pipeline.py +++ b/tests/integration/test_plan_web_order_pipeline.py @@ -67,7 +67,7 @@ def send(url: str, method: str, params: Dict[str, str]) -> str: patch("cyberai.agents.recon.agent.run_nmap", return_value=RECON_RESULT), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), patch("cyberai.agents.recon.agent.discover_surface", return_value=SURFACE), patch("cyberai.agents.exploit.web_exploit._default_sender", return_value=send), diff --git a/tests/unit/test_behavioral_probe.py b/tests/unit/test_behavioral_probe.py index be45d2e..c286fbb 100644 --- a/tests/unit/test_behavioral_probe.py +++ b/tests/unit/test_behavioral_probe.py @@ -128,7 +128,7 @@ def _run_recon(config): patch("cyberai.agents.recon.agent.run_nmap", return_value=ok_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), patch( "cyberai.agents.recon.agent.build_probe_context", @@ -191,7 +191,7 @@ def fake_build(tgt, ports, mass_open=False): patch("cyberai.agents.recon.agent.run_nmap", return_value=mass_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), patch("cyberai.agents.recon.agent.build_probe_context", side_effect=fake_build), ): diff --git a/tests/unit/test_recon.py b/tests/unit/test_recon.py index 982a312..df067c4 100644 --- a/tests/unit/test_recon.py +++ b/tests/unit/test_recon.py @@ -52,7 +52,7 @@ def test_recon_logs_nmap_failure_visibly(monkeypatch): patch("cyberai.agents.recon.agent.run_nmap", return_value=err_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), ): agent.run("t.local") @@ -84,10 +84,52 @@ def test_recon_logs_nmap_success(monkeypatch): patch("cyberai.agents.recon.agent.run_nmap", return_value=ok_nmap), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), ): agent.run("t.local") assert "nmap_scan complete" in logged assert not any("nmap_scan FAILED" in m for m in logged) + + +class TestSubdomainShapeParity: + """The sync ReconAgent and AsyncOrchestrator must write the same shape + under recon.subdomains and derive ReconResult.subdomains identically. + They used to call different enumerators returning different keys, so a + consumer of the KB key was only ever correct for one of the two paths.""" + + _ENUM_RESULT = { + "domain": "t.local", + "found": [{"fqdn": "api.t.local", "ips": ["1.2.3.4"]}], + "count": 1, + "checked": 1, + "wordlist": ["api"], + } + + def test_sync_recon_writes_found_shape_and_fills_recon_result(self): + from unittest.mock import MagicMock, patch + + from cyberai.agents.recon.agent import ReconAgent + from cyberai.core.config import CyberAIConfig + from cyberai.core.scan_session import ScanSession + + session = ScanSession(target="t.local") + agent = ReconAgent(CyberAIConfig(), session, MagicMock(), MagicMock()) + with ( + patch( + "cyberai.agents.recon.agent.run_nmap", + return_value={"target": "t.local", "ports": []}, + ), + patch("cyberai.agents.recon.agent.run_whois", return_value={}), + patch("cyberai.agents.recon.agent.run_dns", return_value={}), + patch( + "cyberai.agents.recon.agent.enumerate_subdomains", + return_value=self._ENUM_RESULT, + ), + patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), + ): + agent.run("t.local") + + assert session.kb.get("recon.subdomains") == self._ENUM_RESULT + assert session.kb.get("recon.result")["subdomains"] == ["api.t.local"] diff --git a/tests/unit/test_recon_guards.py b/tests/unit/test_recon_guards.py index ee2130f..5a37ea3 100644 --- a/tests/unit/test_recon_guards.py +++ b/tests/unit/test_recon_guards.py @@ -14,7 +14,7 @@ _RECON_TOOLS = { "run_whois": {}, "run_dns": {}, - "detect_subdomains": {}, + "enumerate_subdomains": {}, "detect_llm_endpoints": {}, } diff --git a/tests/unit/test_subdomain_enum.py b/tests/unit/test_subdomain_enum.py index 9e73734..f4cae77 100644 --- a/tests/unit/test_subdomain_enum.py +++ b/tests/unit/test_subdomain_enum.py @@ -1,6 +1,7 @@ from unittest.mock import patch from cyberai.agents.recon.subdomain_enum import ( enumerate_subdomains, + fqdns, _resolve, load_wordlist, DEFAULT_WORDLIST, @@ -81,3 +82,34 @@ def test_load_wordlist_from_file(tmp_path): wl.write_text("sub1\nsub2\n# comment\nsub3\n") result = load_wordlist(str(wl)) assert result == ["sub1", "sub2", "sub3"] + + +class TestFqdnsContract: + """fqdns() is the single place that turns an enumerator result into bare + hostnames. Both the sync ReconAgent and AsyncOrchestrator feed + ReconResult.subdomains through it, so a shape change cannot make one + pipeline path silently disagree with the other.""" + + def test_extracts_hostnames_from_found(self): + result = {"found": [{"fqdn": "api.t.local"}, {"fqdn": "dev.t.local"}]} + assert fqdns(result) == ["api.t.local", "dev.t.local"] + + def test_skips_malformed_entries(self): + result = {"found": [{"fqdn": "api.t.local"}, {"ips": ["1.2.3.4"]}, "www.t.local", {}]} + assert fqdns(result) == ["api.t.local"] + + def test_empty_inputs_yield_empty_list(self): + assert fqdns(None) == [] + assert fqdns({}) == [] + assert fqdns({"found": []}) == [] + + def test_real_enumerator_output_is_consumable(self): + """Guards the contract against the producer, not a hand-built dict.""" + from unittest.mock import patch + + with patch( + "cyberai.agents.recon.subdomain_enum._resolve", + side_effect=lambda fqdn: {"fqdn": fqdn, "ips": ["1.2.3.4"]}, + ): + produced = enumerate_subdomains("t.local", wordlist=["api", "dev"]) + assert sorted(fqdns(produced)) == ["api.t.local", "dev.t.local"] diff --git a/tests/unit/test_web_wiring.py b/tests/unit/test_web_wiring.py index cbfd5bc..1d436de 100644 --- a/tests/unit/test_web_wiring.py +++ b/tests/unit/test_web_wiring.py @@ -34,7 +34,7 @@ "run_nmap": {"target": "t.local", "ports": []}, "run_whois": {}, "run_dns": {}, - "detect_subdomains": {}, + "enumerate_subdomains": {}, "detect_llm_endpoints": {}, } @@ -51,7 +51,7 @@ def recon_patches(): patch("cyberai.agents.recon.agent.run_nmap", return_value=_RECON_PATCHES["run_nmap"]), patch("cyberai.agents.recon.agent.run_whois", return_value={}), patch("cyberai.agents.recon.agent.run_dns", return_value={}), - patch("cyberai.agents.recon.agent.detect_subdomains", return_value={}), + patch("cyberai.agents.recon.agent.enumerate_subdomains", return_value={}), patch("cyberai.agents.recon.agent.detect_llm_endpoints", return_value={}), ): yield From 51e542f12554c8fe18656016f61eda1b518ce984 Mon Sep 17 00:00:00 2001 From: Evgeny Kiriyak <224408464+evkir@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:59:46 +0200 Subject: [PATCH 2/4] fix(recon): run whois on the async recon path AsyncReconAgent runs only nmap, dns, subdomains and tls, so recon.whois was never written and ReconResult.whois stayed empty whenever the pipeline went through AsyncOrchestrator, while the sync ReconAgent filled both. Offload the blocking lookup with asyncio.to_thread the same way the LLM endpoint detector is already handled, and feed the result into ReconResult. --- cyberai/core/orchestrator.py | 8 +++++++ tests/integration/test_async_pipeline.py | 29 ++++++++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/cyberai/core/orchestrator.py b/cyberai/core/orchestrator.py index 48c24a8..8944d2b 100644 --- a/cyberai/core/orchestrator.py +++ b/cyberai/core/orchestrator.py @@ -400,6 +400,7 @@ async def _dispatch_async(self, session: ScanSession, phase: ScanPhase) -> Dict[ async def _run_recon_async(self, session: ScanSession) -> Dict: from cyberai.agents.recon.async_agent import AsyncReconAgent + from cyberai.agents.recon.dns_tool import run_whois from cyberai.agents.recon.llm_detector import detect_llm_endpoints from cyberai.agents.recon.subdomain_enum import fqdns from cyberai.core.types import OpenPort, ReconResult @@ -419,6 +420,12 @@ async def _run_recon_async(self, session: ScanSession) -> Dict: llm_result = await asyncio.to_thread(detect_llm_endpoints, session.target) session.kb.set("recon.llm_endpoints", llm_result, agent="async_recon") + # whois too: AsyncReconAgent runs only nmap/dns/subdomains/tls, so the + # sync agent's recon.whois key and ReconResult.whois were empty on the + # async path. Blocking lookup -> offload, same as the LLM detector. + whois_result = await asyncio.to_thread(run_whois, session.target) + session.kb.set("recon.whois", whois_result, agent="async_recon") + # Validated ReconResult so the planner KB graph gets port/service/ # subdomain nodes (build_kb_graph reads recon.result), matching sync. nmap = result.get("nmap") if isinstance(result.get("nmap"), dict) else {} @@ -428,6 +435,7 @@ async def _run_recon_async(self, session: ScanSession) -> Dict: recon_result = ReconResult( target=session.target, ports=[OpenPort(**p) for p in raw_ports if isinstance(p, dict)], + whois=whois_result if isinstance(whois_result, dict) else {}, dns=result.get("dns") if isinstance(result.get("dns"), dict) else {}, subdomains=subdomains, ) diff --git a/tests/integration/test_async_pipeline.py b/tests/integration/test_async_pipeline.py index f08cc98..b8bd0b5 100644 --- a/tests/integration/test_async_pipeline.py +++ b/tests/integration/test_async_pipeline.py @@ -255,6 +255,35 @@ def test_async_recon_populates_result_and_llm_for_planner_graph(self): ntypes = {d["ntype"] for _, d in graph.nodes(data=True)} assert "port" in ntypes and "llm_endpoint" in ntypes + def test_async_recon_writes_whois_like_sync(self): + """AsyncReconAgent runs only nmap/dns/subdomains/tls, so whois has to + be offloaded by the orchestrator or recon.whois and ReconResult.whois + stay empty on the async path while the sync agent fills both.""" + recon_payload = {"target": "t.local", "nmap": {"ports": []}, "dns": {}} + whois_payload = {"target": "t.local", "registrar": "Example Registrar"} + from cyberai.core.config import CyberAIConfig + from cyberai.core.orchestrator import AsyncOrchestrator + from cyberai.core.scan_session import ScanSession + + orch = AsyncOrchestrator(config=CyberAIConfig(), dry_run=False) + session = ScanSession(target="t.local") + with ( + patch( + "cyberai.agents.recon.async_agent.AsyncReconAgent.run", + new_callable=AsyncMock, + return_value=recon_payload, + ), + patch( + "cyberai.agents.recon.llm_detector.detect_llm_endpoints", + return_value={}, + ), + patch("cyberai.agents.recon.dns_tool.run_whois", return_value=whois_payload), + ): + asyncio.run(orch._run_recon_async(session)) + + assert session.kb.get("recon.whois") == whois_payload + assert session.kb.get("recon.result")["whois"] == whois_payload + def test_sync_intel_runs_under_to_thread(self): """Intel/exploit/report stay sync; AsyncOrchestrator must offload them.""" from cyberai.core.scan_session import ScanSession, ScanPhase From 8bc7c021bc9ee5e30067a63d47f4745d5646665e Mon Sep 17 00:00:00 2001 From: Evgeny Kiriyak <224408464+evkir@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:55:46 +0200 Subject: [PATCH 3/4] feat(recon): run web surface discovery on the async recon path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AsyncReconAgent has no web branch, so recon.web_surface was never written under AsyncOrchestrator. Both of its consumers — build_kb_graph, which turns endpoints into HTTP_ENDPOINT nodes for the planner, and ExploitAgent, which reads the key to pick injection points — saw an empty surface, so the entire web layer was dead on the async path while working on the sync one. Reuse ReconAgent._run_web_recon through asyncio.to_thread instead of re-deriving the crawl and its findings in the orchestrator, gated by the same use_web_recon flag as the sync path. --- cyberai/core/orchestrator.py | 11 +++++ tests/integration/test_async_pipeline.py | 54 ++++++++++++++++++++++++ 2 files changed, 65 insertions(+) diff --git a/cyberai/core/orchestrator.py b/cyberai/core/orchestrator.py index 8944d2b..4c44bbe 100644 --- a/cyberai/core/orchestrator.py +++ b/cyberai/core/orchestrator.py @@ -426,6 +426,17 @@ async def _run_recon_async(self, session: ScanSession) -> Dict: whois_result = await asyncio.to_thread(run_whois, session.target) session.kb.set("recon.whois", whois_result, agent="async_recon") + # HTTP attack surface: AsyncReconAgent has no web branch, so + # recon.web_surface was absent on the async path and every consumer + # of it (build_kb_graph, ExploitAgent) silently saw nothing. Reuse the + # sync agent's method rather than re-deriving the crawl and its + # findings here — one source of truth for the web surface. + if getattr(self.config, "use_web_recon", False): + from cyberai.agents.recon.agent import ReconAgent + + web_agent = ReconAgent(self.config, session, None, getattr(self, "audit", None)) + await asyncio.to_thread(web_agent._run_web_recon, session.target) + # Validated ReconResult so the planner KB graph gets port/service/ # subdomain nodes (build_kb_graph reads recon.result), matching sync. nmap = result.get("nmap") if isinstance(result.get("nmap"), dict) else {} diff --git a/tests/integration/test_async_pipeline.py b/tests/integration/test_async_pipeline.py index b8bd0b5..c99e8cb 100644 --- a/tests/integration/test_async_pipeline.py +++ b/tests/integration/test_async_pipeline.py @@ -284,6 +284,60 @@ def test_async_recon_writes_whois_like_sync(self): assert session.kb.get("recon.whois") == whois_payload assert session.kb.get("recon.result")["whois"] == whois_payload + def test_async_recon_runs_web_surface_when_flag_is_on(self): + """AsyncReconAgent has no web branch. Without the orchestrator running + one, recon.web_surface is absent on the async path and build_kb_graph + plus ExploitAgent see an empty surface — the whole web layer is dead + under AsyncOrchestrator while it works under the sync one.""" + from cyberai.core.config import CyberAIConfig + from cyberai.core.orchestrator import AsyncOrchestrator + from cyberai.core.scan_session import ScanSession + + recon_payload = {"target": "t.local", "nmap": {"ports": []}, "dns": {}} + surface = { + "endpoints": [{"method": "GET", "url": "http://t.local/search", "params": ["q"]}] + } + orch = AsyncOrchestrator(config=CyberAIConfig(use_web_recon=True), dry_run=False) + session = ScanSession(target="t.local") + with ( + patch( + "cyberai.agents.recon.async_agent.AsyncReconAgent.run", + new_callable=AsyncMock, + return_value=recon_payload, + ), + patch("cyberai.agents.recon.llm_detector.detect_llm_endpoints", return_value={}), + patch("cyberai.agents.recon.dns_tool.run_whois", return_value={}), + patch("cyberai.agents.recon.agent.discover_surface", return_value=surface), + ): + asyncio.run(orch._run_recon_async(session)) + + assert session.kb.get("recon.web_surface") == surface + + def test_async_recon_skips_web_surface_when_flag_is_off(self): + """Default config must not start crawling: the flag gates the async + path exactly as it gates the sync one.""" + from cyberai.core.config import CyberAIConfig + from cyberai.core.orchestrator import AsyncOrchestrator + from cyberai.core.scan_session import ScanSession + + recon_payload = {"target": "t.local", "nmap": {"ports": []}, "dns": {}} + orch = AsyncOrchestrator(config=CyberAIConfig(), dry_run=False) + session = ScanSession(target="t.local") + with ( + patch( + "cyberai.agents.recon.async_agent.AsyncReconAgent.run", + new_callable=AsyncMock, + return_value=recon_payload, + ), + patch("cyberai.agents.recon.llm_detector.detect_llm_endpoints", return_value={}), + patch("cyberai.agents.recon.dns_tool.run_whois", return_value={}), + patch("cyberai.agents.recon.agent.discover_surface") as crawl, + ): + asyncio.run(orch._run_recon_async(session)) + + crawl.assert_not_called() + assert session.kb.get("recon.web_surface") is None + def test_sync_intel_runs_under_to_thread(self): """Intel/exploit/report stay sync; AsyncOrchestrator must offload them.""" from cyberai.core.scan_session import ScanSession, ScanPhase From 7f834db0fde49dc3c13bfa16459cc707229009e5 Mon Sep 17 00:00:00 2001 From: Evgeny Kiriyak <224408464+evkir@users.noreply.github.com> Date: Thu, 30 Jul 2026 21:59:11 +0200 Subject: [PATCH 4/4] fix(llm): build the ollama request in one place for sync and async _acall_ollama constructed its own payload and had drifted from the sync one: it dropped the system prompt entirely, omitted the raised num_ctx that keeps long exploit prompts inside the context window, and allowed half the timeout. Any agent migrated to acall would have lost its system prompt silently. Extract _ollama_request as the single builder both paths call and hoist the timeout into OLLAMA_TIMEOUT, so the two cannot diverge again. --- cyberai/core/llm_client.py | 33 +++++++++++++------- tests/unit/test_llm_client.py | 58 +++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 11 deletions(-) diff --git a/cyberai/core/llm_client.py b/cyberai/core/llm_client.py index fac4322..3471757 100644 --- a/cyberai/core/llm_client.py +++ b/cyberai/core/llm_client.py @@ -10,6 +10,11 @@ from .base_agent import Tool +# Local models on a 3060 are slow on long exploit prompts; the async path +# used to allow only 60s and timed out where the sync path succeeded. +OLLAMA_TIMEOUT = 120 + + class LLMClient: """ Unified LLM interface — OpenAI / Anthropic / Ollama @@ -122,9 +127,14 @@ def _call_anthropic( ) return response.content[0].text - def _call_ollama( - self, messages: List[Dict], system: Optional[str], agent_name: str = "unknown" - ) -> str: + def _ollama_request(self, messages: List[Dict], system: Optional[str]) -> tuple: + """Build the (url, payload) pair for an ollama /api/chat call. + + Single source for both the sync and async entry points: the async one + used to build its own payload and had silently dropped the system + prompt and the raised num_ctx, so the same call behaved differently + depending on which path reached it. + """ url = f"{self.config.base_url or 'http://localhost:11434'}/api/chat" full_messages: List[Dict] = [] if system: @@ -138,7 +148,13 @@ def _call_ollama( # prompts (CVE JSON + attack paths + chain), which 4xx/5xx the call. "options": {"num_ctx": 8192}, } - response = httpx.post(url, json=payload, timeout=120) + return url, payload + + def _call_ollama( + self, messages: List[Dict], system: Optional[str], agent_name: str = "unknown" + ) -> str: + url, payload = self._ollama_request(messages, system) + response = httpx.post(url, json=payload, timeout=OLLAMA_TIMEOUT) if response.status_code != 200: raise RuntimeError(f"ollama HTTP {response.status_code}: {response.text[:300]}") data = response.json() @@ -420,13 +436,8 @@ async def _acall_anthropic( async def _acall_ollama( self, messages: List[Dict], system: Optional[str], agent_name: str = "unknown" ) -> str: - url = f"{self.config.base_url or 'http://localhost:11434'}/api/chat" - payload = { - "model": self.config.model, - "messages": messages, - "stream": False, - } - async with httpx.AsyncClient(timeout=60) as client: + url, payload = self._ollama_request(messages, system) + async with httpx.AsyncClient(timeout=OLLAMA_TIMEOUT) as client: response = await client.post(url, json=payload) response.raise_for_status() data = response.json() diff --git a/tests/unit/test_llm_client.py b/tests/unit/test_llm_client.py index 183fbbb..a441cbc 100644 --- a/tests/unit/test_llm_client.py +++ b/tests/unit/test_llm_client.py @@ -145,3 +145,61 @@ async def post(self, url, json): assert client.cost_tracker.call_count == 1 assert client.cost_tracker.total_input_tokens == 10 assert client.cost_tracker.total_output_tokens == 3 + + +def test_ollama_sync_and_async_send_identical_payloads(monkeypatch): + """Both ollama entry points must build the same request body. The async + one used to construct its own, dropping the system prompt and the raised + num_ctx, so the same prompt behaved differently depending on the path.""" + import asyncio + + from cyberai.core import llm_client as lc + + captured = {} + + class _Resp: + status_code = 200 + + def raise_for_status(self): + pass + + def json(self): + return {"message": {"content": "ok"}} + + def fake_post(url, json, timeout): + captured["sync"] = (url, json, timeout) + return _Resp() + + class _AsyncClient: + def __init__(self, timeout): + captured["async_timeout"] = timeout + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def post(self, url, json): + captured["async"] = (url, json) + return _Resp() + + monkeypatch.setattr(lc.httpx, "post", fake_post) + monkeypatch.setattr(lc.httpx, "AsyncClient", _AsyncClient) + + client = lc.LLMClient.__new__(lc.LLMClient) + client.config = type("C", (), {"base_url": None, "model": "qwen2.5:7b", "provider": "ollama"})() + client.cost_tracker = None + client.budget_usd = 0 + + messages = [{"role": "user", "content": "hi"}] + client._call_ollama(messages, "SYS") + asyncio.run(client._acall_ollama(messages, "SYS")) + + sync_url, sync_payload, sync_timeout = captured["sync"] + async_url, async_payload = captured["async"] + assert sync_url == async_url + assert sync_payload == async_payload + assert sync_timeout == captured["async_timeout"] == lc.OLLAMA_TIMEOUT + assert async_payload["messages"][0] == {"role": "system", "content": "SYS"} + assert async_payload["options"]["num_ctx"] == 8192