From 422f68141516eb942d837974ec942c3fe3bd1dff Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Thu, 26 Mar 2026 11:18:27 +0100 Subject: [PATCH 01/12] Traversale is a protocol, to it can be of instance of Path --- .../src/deadend_agent/embedders/code_indexer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/embedders/code_indexer.py b/deadend_cli/deadend_agent/src/deadend_agent/embedders/code_indexer.py index 504087f..f0efd88 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/embedders/code_indexer.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/embedders/code_indexer.py @@ -302,8 +302,9 @@ def _load_patterns(self): # Get the path to the data file within the package data_file = importlib.resources.files('deadend_agent').joinpath('data/vendor_specific_files.json') - with open(data_file, encoding="utf-8") as f: - patterns_json = f.read() + if isinstance(data_file, Path): + with open(data_file, encoding="utf-8") as f: + patterns_json = f.read() self.forbidden_patterns = json.loads(patterns_json) From dae623601d2318be6dd9ccced7e0afcf2330ef1f Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Fri, 27 Mar 2026 19:10:05 +0100 Subject: [PATCH 02/12] adding Agent memory setup. Moving the setup directory to environments. First setup of an agentic Virtual filesystem, with read,write,ripgrep and list to gather information from memory. I've also changing the realm where the agent operates. Now everything goes to .cache/deadend/agents/... and saves the agent id used in config so that we can gather the agent information directly. the memory files are present inside that place. and they are simple MD files gathered after each agent run. So that the next agent knows what's going on. --- .gitmodules | 3 + benchmarks/xbow/validation-benchmarks-xbow | 2 +- deadend_cli/deadend_agent/pyproject.toml | 2 + .../src/deadend_agent/agents/__init__.py | 6 +- .../agents/components/executor.py | 271 +++---- .../deadend_agent/agents/exploit_web_agent.py | 5 +- .../agents/generic_agents/__init__.py | 4 +- .../agents/generic_agents/memory_agent.py | 64 ++ .../python_interpreter_agent.py | 2 +- .../agents/generic_agents/request_agent.py | 2 +- .../agents/generic_agents/shell_agent.py | 4 +- .../generic_agents/webapp_analyzer_agent.py | 6 +- .../agents/recon_threatmodel_agent.py | 4 +- .../src/deadend_agent/config/settings.py | 8 +- .../src/deadend_agent/context/__init__.py | 4 +- .../deadend_agent/context/context_engine.py | 53 +- .../src/deadend_agent/context/memory.py | 177 ++--- .../src/deadend_agent/core_agent/__init__.py | 3 + .../deadend_agent/core_agent/rlm_runner.py | 716 ++++++++++++++++++ .../src/deadend_agent/deadend_agent.py | 108 ++- .../src/deadend_agent/rlm/__init__.py | 16 + .../src/deadend_agent/rlm/compat.py | 52 ++ .../src/deadend_agent/rlm/memory.py | 473 ++++++++++++ .../src/deadend_agent/tools/__init__.py | 11 +- .../src/deadend_agent/tools/avfs/__init__.py | 17 + .../src/deadend_agent/tools/avfs/avfs.py | 183 +++++ .../src/deadend_agent/tools/avfs/list.py | 120 +++ .../src/deadend_agent/tools/avfs/read.py | 157 ++++ .../src/deadend_agent/tools/avfs/write.py | 53 ++ .../tools/python_interpreter/__init__.py | 11 +- .../python_interpreter/python_interpreter.py | 199 +---- .../src/deadend_agent/utils/structures.py | 18 + .../deadend_agent/tests/rlm/test_avfs.py | 254 +++++++ .../rlm/test_deadend_agent_avfs_startup.py | 450 +++++++++++ .../deadend_agent/tests/rlm/test_memory.py | 100 +++ .../tests/rlm/test_memory_avfs.py | 72 ++ .../deadend_agent/tests/rlm/test_runner.py | 190 +++++ .../deadend_eval/src/deadend_eval/eval.py | 10 +- .../_shared/_memory_summary.jinja2 | 21 + .../memory.instructions.jinja2 | 179 +++++ .../deadend_prompts/shell.instructions.jinja2 | 13 +- .../src/deadend_prompts/template_renderer.py | 2 +- .../tools/avfs_grep.description.jinja2 | 1 + .../tools/avfs_list.description.jinja2 | 1 + .../tools/avfs_read.description.jinja2 | 1 + .../tools/avfs_write.description.jinja2 | 1 + deadend_cli/pyproject.toml | 5 +- deadend_cli/simple-python-interpreter-sandbox | 1 + deadend_cli/src/deadend_cli/chat.py | 15 +- deadend_cli/src/deadend_cli/cli.py | 2 + deadend_cli/src/deadend_cli/jsonrpc_server.py | 7 +- deadend_cli/uv.lock | 18 + docs/RLM.md | 359 +++++++++ docs/RLM_as_memory.md | 663 ++++++++++++++++ {setup => environments}/gvisor/daemon.json | 0 .../gvisor/daemon.json.new | 0 .../gvisor/install_gvisor.sh | 0 .../images/kalilinux.Dockerfile | 0 .../images/webapp_sec.Dockerfile | 0 .../pgvector/setup_pgvector.sh | 0 60 files changed, 4637 insertions(+), 482 deletions(-) create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/core_agent/rlm_runner.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/rlm/__init__.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/rlm/compat.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/rlm/memory.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/avfs.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py create mode 100644 deadend_cli/deadend_agent/tests/rlm/test_avfs.py create mode 100644 deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py create mode 100644 deadend_cli/deadend_agent/tests/rlm/test_memory.py create mode 100644 deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py create mode 100644 deadend_cli/deadend_agent/tests/rlm/test_runner.py create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_grep.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_list.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_read.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_write.description.jinja2 create mode 160000 deadend_cli/simple-python-interpreter-sandbox create mode 100644 docs/RLM.md create mode 100644 docs/RLM_as_memory.md rename {setup => environments}/gvisor/daemon.json (100%) rename {setup => environments}/gvisor/daemon.json.new (100%) rename {setup => environments}/gvisor/install_gvisor.sh (100%) rename {setup => environments}/images/kalilinux.Dockerfile (100%) rename {setup => environments}/images/webapp_sec.Dockerfile (100%) rename {setup => environments}/pgvector/setup_pgvector.sh (100%) diff --git a/.gitmodules b/.gitmodules index d26c8e7..dcf0f99 100644 --- a/.gitmodules +++ b/.gitmodules @@ -4,3 +4,6 @@ [submodule "benchmarks/xbow/validation-benchmarks-xbow"] path = benchmarks/xbow/validation-benchmarks-xbow url = https://github.com/xoxruns/validation-benchmarks-xbow.git +[submodule "deadend_cli/simple-python-interpreter-sandbox"] + path = deadend_cli/simple-python-interpreter-sandbox + url = https://github.com/xoxruns/simple-python-interpreter-sandbox.git diff --git a/benchmarks/xbow/validation-benchmarks-xbow b/benchmarks/xbow/validation-benchmarks-xbow index f5e1a36..320af47 160000 --- a/benchmarks/xbow/validation-benchmarks-xbow +++ b/benchmarks/xbow/validation-benchmarks-xbow @@ -1 +1 @@ -Subproject commit f5e1a36a5bc1938e79076460cf0a5b58dc22af6f +Subproject commit 320af47d848b94370655228b61e0c62443de15b0 diff --git a/deadend_cli/deadend_agent/pyproject.toml b/deadend_cli/deadend_agent/pyproject.toml index b89de2c..eb43cb3 100644 --- a/deadend_cli/deadend_agent/pyproject.toml +++ b/deadend_cli/deadend_agent/pyproject.toml @@ -29,12 +29,14 @@ dependencies = [ "opentelemetry-exporter-otlp>=1.39.1", "opentelemetry-sdk>=1.39.1", "playwright>=1.56.0", + "python-sandbox-client", "prompt-toolkit>=3.0.51", "pydantic>=2.11.5", "pydantic-ai>=1.35.0", "pydantic-ai-slim[google,openrouter]>=1.35.0", "pyyaml>=6.0.2", "readchar>=4.2.1", + "ripgrepy>=2.2.0", "rich>=14.0.0", "semantic-text-splitter>=0.27.0", "sqlalchemy>=2.0.41", diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py index 8ea4c1d..efecc1f 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py @@ -10,6 +10,7 @@ from .generic_agents.python_interpreter_agent import PythonInterpreterAgent, PythonInterpreterOutput from .generic_agents.request_agent import RequesterAgent, RequesterOutput from .generic_agents.webapp_analyzer_agent import WebAppAnalyzerAgent +from .generic_agents.memory_agent import MemoryAgent __all__ = [ "AgentRunner", "AgentOutput", @@ -20,6 +21,7 @@ "ShellAgent", "ShellOutput", "PythonInterpreterAgent", "PythonInterpreterOutput", "RequesterAgent", "RequesterOutput", - "WebAppAnalyzerAgent" + "WebAppAnalyzerAgent", + "MemoryAgent" -] \ No newline at end of file +] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py index 56bee8e..21c92ac 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py @@ -8,12 +8,14 @@ RequesterAgent, ShellAgent, PythonInterpreterAgent, AgentOutput, - WebAppAnalyzerAgent + WebAppAnalyzerAgent, + MemoryAgent, ) from deadend_agent.agents.components.planner import TaskNode from deadend_agent.context import ContextEngine from deadend_agent.config.settings import ModelSpec -from deadend_agent.utils.structures import WebappreconDeps, RequesterDeps, ShellDeps +from deadend_agent.tools.avfs.write import write_text +from deadend_agent.utils.structures import MemoryWorkspaceDeps, WebappreconDeps, RequesterDeps, ShellDeps class LogEvent(BaseModel): @@ -31,6 +33,29 @@ class ResultEvent(BaseModel): # Union type for all possible executor events ExecutorEvent = LogEvent | ResultEvent + +def _memory_prompt_prefix(memory_context: str) -> str: + """Render persistent memory context as a prompt prefix for downstream agents.""" + if not memory_context.strip(): + return "" + return f"## Persistent Memory Context\n{memory_context.strip()}\n\n" + + +def _build_memory_summary(agent_name: str, task: str, output: AgentOutput) -> str: + """Create a deterministic memory entry from structured agent output.""" + summary = output.detailed_summary.strip() or "None" + proofs = output.proofs.strip() or "None" + thoughts = output.thoughts.strip() or "None" + return ( + "## Task Summary\n" + f"- Agent: {agent_name}\n" + f"- Task: {task.strip()}\n" + f"- Confidence: {output.confidence_score:.2f}\n" + f"- Summary: {summary}\n" + f"- Proofs: {proofs}\n" + f"- Thoughts: {thoughts}\n\n" + ) + @dataclass class SupervisorDeps: """Dependencies for supervisor router containing all agents and their deps.""" @@ -40,10 +65,14 @@ class SupervisorDeps: shell_deps: ShellDeps | None python_interpreter_agent: PythonInterpreterAgent webapp_analyzer_agent: WebAppAnalyzerAgent + memory_agent: MemoryAgent + memory_deps: MemoryWorkspaceDeps session_id: str message_history: list | None usage_limits: UsageLimits deferred_tool_results: DeferredToolResults | None + memory_context: str = "" + auth_session_key: str = "" context: ContextEngine | None = None # Context engine for storing agent outputs class AgentExecutor: @@ -80,6 +109,8 @@ def __init__( self.requires_approval = requires_approval self.context = context self.session_id = session_id + self.memory_context = "" + self.auth_session_key = "" self.supervisor = SupervisorAgent( model=self.model, @@ -106,124 +137,56 @@ def set_dependencies( if webapprecon_deps is not None: self.webapprecon_deps = webapprecon_deps - # def _executor_message_yield(self, message) -> SupervisorOutput | AgentOutput | LogEvent | ResultEvent: - # pass - - # async def execute( - # self, - # task_node: TaskNode, - # agent_context: str = "", - # usage: RunUsage = RunUsage(), - # usage_limits: UsageLimits = UsageLimits(request_limit=None, tool_calls_limit=None), - # deferred_tool_results: DeferredToolResults | None = None, - # message_history: list | None = None - # ) -> AsyncGenerator[ExecutorEvent, None]: - # """Execute a task node using the appropriate agent. - - # The execution process: - # 1. Uses the router (if available) to determine which agent should handle the task - # 2. Attempts to get or create the selected specialized agent - # 3. Executes the task with the specialized agent or falls back to the generic runner - # 4. Extracts confidence score and updates context with execution results - - # Args: - # task_node: The TaskNode containing the task to execute - # context: Current execution context (will be copied and updated) - # deps: Optional dependencies to pass to the agent - # usage: Usage tracking object - # usage_limits: Limits for token usage - # deferred_tool_results: Optional deferred tool results from previous runs - # message_history: Previous conversation messages for context - - # Yields: - # LogEvent instances for streaming updates. - # The final event is a ResultEvent instance. - - # Note: - # If routing fails or the selected agent cannot be created, execution falls back - # to the generic runner. All routing and execution information is logged in the context. - # """ - # context: dict[str, Any] = {"log": ""} - # confidence_score: float | None = None - - # def emit(message: str) -> LogEvent: - # """Append a log entry to the context and return it for streaming.""" - # context["log"] += f"\n{message}" - # return LogEvent(message=message) - - # try: - # yield emit(f"Current task: {task_node.task}\n") - - # routing_info = None - # selected_agent: AgentRunner | None = None - # if self.router: - # try: - # router_result = await self.router.run( - # prompt=f"{agent_context}\nWhich agent should handle: {task_node.task}", - # deps=None, - # message_history=message_history or "", - # usage=usage, - # usage_limits=usage_limits, - # deferred_tool_results=None - # ) - # routing_info = router_result.output - # if isinstance(routing_info, RouterOutput): - # yield emit( - # "Selected agent: " - # f"{routing_info.next_agent_name}\nReasoning: {routing_info.reasoning}" - # ) - # selected_agent = self._get_agent(routing_info.next_agent_name) - # if selected_agent: - # yield emit(f"Using specialized agent: {routing_info.next_agent_name}") - # except Exception as exc: - # yield emit(f"Routing failed: {exc}, using generic executor") - - # if isinstance(selected_agent, AgentRunner): - # result = await self._run_agent( - # agent=selected_agent, - # prompt=agent_context+task_node.task, - # message_history=message_history, - # usage=usage, - # usage_limits=usage_limits, - # deferred_tool_results=deferred_tool_results - # ) - # output = result.output - # # print(f"test output : {output}") - # else: - # output = f"[AGENT RESPONSE] Error in agent running {selected_agent}" - - # notes = "" - # updated_state = {} - # if isinstance(output, AgentOutput): - # confidence_score = output.confidence_score - # notes = output.notes - # updated_state = output.updated_state or {} - # else: - # # Default confidence score when output is not an AgentOutput - # confidence_score = 0.5 - - # # yield emit(f"[EXECUTOR] Task: {task_node.task}\nNotes: {notes}\n{output}") - # context.update(updated_state) - # context["last_output"] = output.model_dump() if isinstance(output, AgentOutput) else str(output) - # yield ResultEvent( - # confidence_score=confidence_score, - # context=context, - # ) - # return - # except UsageLimitExceeded as exc: - # yield emit(f"[EXECUTOR] Usage limit reached: {exc}") - # yield ResultEvent( - # confidence_score=confidence_score or 0.5, - # context=context, - # ) - # return - # except Exception as exc: - # yield emit(f"[EXECUTOR] Error: {exc}") - # yield ResultEvent( - # confidence_score=confidence_score or 0.5, - # context=context, - # ) - # return + def set_memory_context(self, memory_context: str) -> None: + """Register startup memory context for downstream agents.""" + self.memory_context = memory_context + + def set_auth_session_key(self, auth_session_key: str) -> None: + """Register the auth storage session key used by the python interpreter agent.""" + self.auth_session_key = auth_session_key + + async def _refresh_memory_context_for_task(self, task_query: str) -> str: + """Retrieve task-specific memory immediately before supervisor execution.""" + memory_workspace_root = ( + self.requester_deps.memory_workspace_root + if self.requester_deps is not None + else (self.shell_deps.memory_workspace_root if self.shell_deps is not None else None) + ) + if memory_workspace_root is None or self.session_id is None: + self.memory_context = "" + return self.memory_context + + memory_agent = MemoryAgent( + model=self.model, + deps_type=MemoryWorkspaceDeps, + ) + memory_deps = MemoryWorkspaceDeps( + session_id=self.session_id, + memory_workspace_root=memory_workspace_root, + ) + result = await memory_agent.run( + prompt=( + f"Current task:\n{task_query}\n\n" + "Inspect the persistent memory workspace using AVFS tools with workspace=\"memory\". " + "Return only a concise task-relevant memory summary as plain text for the supervisor. " + "If memory is empty or not useful for this task, return a short plain-text statement saying that no relevant persisted memory is available." + ), + deps=memory_deps, + message_history=[], + usage=RunUsage(), + usage_limits=UsageLimits(request_limit=None, tool_calls_limit=None), + deferred_tool_results=None, + ) + + output = getattr(result, "output", None) + self.memory_context = str(output).strip() if output is not None else "" + if self.shell_deps is not None: + self.shell_deps.memory_context = self.memory_context + if self.requester_deps is not None: + self.requester_deps.memory_context = self.memory_context + if self.webapprecon_deps is not None: + self.webapprecon_deps.memory_context = self.memory_context + return self.memory_context async def execute_supervisor( self, @@ -268,6 +231,7 @@ def emit(message: str) -> LogEvent: return LogEvent(message=message) try: + await self._refresh_memory_context_for_task(task_node.task) yield emit(f"Current task: {task_node.task}\n") # Instantiate all generic agents requester_agent = RequesterAgent( @@ -281,12 +245,12 @@ def emit(message: str) -> LogEvent: model=self.model, deps_type=WebappreconDeps, target_information=self.context.target, - requires_approval=self.requires_approval + requires_approval=self.requires_approval, ) if self.shell_deps is not None else None python_interpreter_agent = PythonInterpreterAgent( model=self.model, - deps_type=str, + deps_type=MemoryWorkspaceDeps, ) webapp_analyzer_agent = WebAppAnalyzerAgent( @@ -294,6 +258,19 @@ def emit(message: str) -> LogEvent: deps_type=RequesterDeps, ) + memory_deps = MemoryWorkspaceDeps( + session_id=self.session_id or "", + memory_workspace_root=( + self.requester_deps.memory_workspace_root + if self.requester_deps is not None + else (self.shell_deps.memory_workspace_root if self.shell_deps is not None else None) + ), + memory_context=self.memory_context, + ) + memory_agent = MemoryAgent( + model=self.model, + deps_type=MemoryWorkspaceDeps, + ) # Create supervisor dependencies supervisor_deps = SupervisorDeps( @@ -303,10 +280,14 @@ def emit(message: str) -> LogEvent: shell_deps=self.shell_deps, python_interpreter_agent=python_interpreter_agent, webapp_analyzer_agent=webapp_analyzer_agent, + memory_agent=memory_agent, + memory_deps=memory_deps, session_id=self.session_id or "", message_history=message_history, usage_limits=usage_limits, deferred_tool_results=deferred_tool_results, + memory_context=self.memory_context, + auth_session_key=self.auth_session_key, context=self.context # Pass context for storing agent outputs ) @@ -380,6 +361,16 @@ def _add_agent_output_to_context( skip_structured=False ) + def _persist_agent_summary(agent_name: str, task: str, output: AgentOutput) -> None: + """Persist a deterministic summary into the memory workspace.""" + write_text( + f"summaries/{agent_name}.md", + _build_memory_summary(agent_name, task, output), + session_id=self.session_id, + workspace="memory", + append=True, + ) + # Create tool functions using RunContext for agent delegation @supervisor.agent.tool async def call_requester_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: @@ -388,8 +379,9 @@ async def call_requester_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> if ctx.deps.requester_agent is None or ctx.deps.requester_deps is None: return "Requester agent dependencies not configured." + memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) result = await ctx.deps.requester_agent.run( - prompt, + f"{memory_prefix}{prompt}", deps=ctx.deps.requester_deps, message_history=ctx.deps.message_history, usage=ctx.usage, @@ -404,6 +396,7 @@ async def call_requester_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> agent_name="requester", output=result.output ) + _persist_agent_summary("requester", prompt, result.output) return f"Requester agent result: {result.output.model_dump()}" return str(result.output) if hasattr(result, 'output') else str(result) @@ -413,8 +406,9 @@ async def call_shell_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: print(f"input tool looking for the error : {prompt}") if ctx.deps.shell_agent is None or ctx.deps.shell_deps is None: return "Shell agent dependencies not configured." + memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) result = await ctx.deps.shell_agent.run( - prompt, + f"{memory_prefix}{prompt}", deps=ctx.deps.shell_deps, message_history=ctx.deps.message_history, usage=ctx.usage, @@ -429,6 +423,7 @@ async def call_shell_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: agent_name="shell", output=result.output ) + _persist_agent_summary("shell", prompt, result.output) return f"Shell agent result: {result.output.model_dump()}" return str(result.output) if hasattr(result, 'output') else str(result) @@ -437,8 +432,9 @@ async def call_webapp_analyzer_agent(ctx: RunContext[SupervisorDeps], prompt: st print(f"input tool looking for the error : {prompt}") print(ctx.deps.requester_deps) + memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) result = await ctx.deps.webapp_analyzer_agent.run( - prompt, + f"{memory_prefix}{prompt}", deps=ctx.deps.requester_deps, message_history=ctx.deps.message_history, usage=ctx.usage, @@ -453,10 +449,11 @@ async def call_python_interpreter_agent(ctx: RunContext[SupervisorDeps], prompt: """Call the python interpreter agent to execute Python scripts.""" print(f"input tool looking for the error : {prompt}") + memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) result = await ctx.deps.python_interpreter_agent.run( - prompt, - deps=ctx.deps.session_id, - session_key=ctx.deps.session_id, + f"{memory_prefix}{prompt}", + deps=ctx.deps.memory_deps, + session_key=ctx.deps.auth_session_key, message_history=ctx.deps.message_history, usage=ctx.usage, usage_limits=ctx.deps.usage_limits, @@ -469,12 +466,28 @@ async def call_python_interpreter_agent(ctx: RunContext[SupervisorDeps], prompt: context=ctx.deps.context, agent_name="python_interpreter", output=result.output - ) + ) + _persist_agent_summary("python_interpreter", prompt, result.output) return f"Python interpreter agent result: {result.output.model_dump()}" return str(result.output) if hasattr(result, 'output') else str(result) + @supervisor.agent.tool + async def call_memory_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: + """Call the memory agent to inspect or update persistent notes.""" + result = await ctx.deps.memory_agent.run( + prompt, + deps=ctx.deps.memory_deps, + message_history=ctx.deps.message_history, + usage=ctx.usage, + usage_limits=ctx.deps.usage_limits, + deferred_tool_results=ctx.deps.deferred_tool_results, + ) + return str(result.output.model_dump()) if hasattr(result, "output") else str(result) + # Execute task with supervisor supervisor_prompt = f"Your task is : {task_node.task}\n" + if supervisor_deps.memory_context: + supervisor_prompt += f"## Persistent Memory Context:\n{supervisor_deps.memory_context}\n" if agent_context: supervisor_prompt += f"## Traces: \n{agent_context}\n" diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/exploit_web_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/exploit_web_agent.py index d41fd46..e766eed 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/exploit_web_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/exploit_web_agent.py @@ -12,7 +12,6 @@ from typing import Any from pydantic import BaseModel -from pydantic_ai import Tool from deadend_agent.config.settings import ModelSpec from deadend_agent.utils.structures import PlannerOutput from deadend_agent.tools import ( @@ -65,7 +64,9 @@ async def run( message_history, usage, usage_limits, - deferred_tool_results=None + deferred_tool_results=None, + *args, + **kwargs ): return await super().run( prompt, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/__init__.py index c19c6a2..3c0b04a 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/__init__.py @@ -1,5 +1,7 @@ +from .memory_agent import MemoryAgent from .python_interpreter_agent import PythonInterpreterAgent from .request_agent import RequesterAgent from .shell_agent import ShellAgent +from .webapp_analyzer_agent import WebAppAnalyzerAgent -__all__ = ["PythonInterpreterAgent", "RequesterAgent", "ShellAgent"] \ No newline at end of file +__all__ = ["MemoryAgent", "PythonInterpreterAgent", "RequesterAgent", "ShellAgent", "WebAppAnalyzerAgent"] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py new file mode 100644 index 0000000..6fdc771 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py @@ -0,0 +1,64 @@ +from typing import Any + +from pydantic_ai import DeferredToolRequests, DeferredToolResults, Tool +from pydantic_ai.usage import RunUsage, UsageLimits + +from deadend_agent.agents.factory import AgentRunner +from deadend_agent.config.settings import ModelSpec +from deadend_agent.tools import avfs_grep, avfs_list, avfs_read, avfs_write +from deadend_prompts import render_agent_instructions, render_tool_description + + +class MemoryAgent(AgentRunner): + """Agent dedicated to reading and writing the persistent memory workspace.""" + + def __init__( + self, + model: ModelSpec, + deps_type: Any | None, + ): + tools_metadata = { + "avfs_list": render_tool_description("avfs_list"), + "avfs_read": render_tool_description("avfs_read"), + "avfs_write": render_tool_description("avfs_write"), + "avfs_grep": render_tool_description("avfs_grep"), + } + + self.instructions = render_agent_instructions( + agent_name="memory", + tools=tools_metadata, + ) + + super().__init__( + name="memory", + model=model, + instructions=self.instructions, + deps_type=deps_type, + output_type=[str, DeferredToolRequests], + tools=[ + Tool(avfs_list), + Tool(avfs_read), + Tool(avfs_write), + Tool(avfs_grep), + ], + ) + + async def run( + self, + prompt, + deps, + message_history, + usage: RunUsage | None, + usage_limits: UsageLimits | None, + deferred_tool_results: DeferredToolResults | None = None, + *args, + **kwargs, + ): + return await super().run( + prompt=prompt, + deps=deps, + message_history=message_history, + usage=usage, + usage_limits=usage_limits, + deferred_tool_results=deferred_tool_results, + ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/python_interpreter_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/python_interpreter_agent.py index ea6e078..c130cf7 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/python_interpreter_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/python_interpreter_agent.py @@ -14,7 +14,7 @@ # from deadend_agent.context import MemoryHandler from deadend_agent.config.settings import ModelSpec from deadend_agent.agents.factory import AgentRunner, AgentOutput -from deadend_agent.tools import run_python_file, read_auth_storage +from deadend_agent.tools import read_auth_storage, run_python_file from deadend_prompts import render_agent_instructions, render_tool_description class PythonInterpreterOutput(AgentOutput): diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/request_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/request_agent.py index 039c5a6..5e7c821 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/request_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/request_agent.py @@ -87,7 +87,7 @@ def __init__( deps_type=deps_type, output_type=[RequesterOutput, DeferredToolRequests], tools=[ - Tool(pw_send_payload, requires_approval=requires_approval) + Tool(pw_send_payload, requires_approval=requires_approval), ] ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/shell_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/shell_agent.py index fe65f2d..1d916ee 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/shell_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/shell_agent.py @@ -55,9 +55,7 @@ def __init__( instructions=self.instructions, deps_type=deps_type, output_type=[ShellOutput, DeferredToolRequests], - tools=[ - Tool(sandboxed_shell_tool, requires_approval=requires_approval), - ] + tools=[Tool(sandboxed_shell_tool, requires_approval=requires_approval)], ) async def run( diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/webapp_analyzer_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/webapp_analyzer_agent.py index 528d04d..d408634 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/webapp_analyzer_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/webapp_analyzer_agent.py @@ -18,7 +18,7 @@ def __init__( deps_type: Any | None, ): tools_metadata = { - "webapp_analyzer": render_tool_description("webapp_analyzer") + "webapp_analyzer": render_tool_description("webapp_analyzer"), } self.instructions = render_agent_instructions( @@ -33,7 +33,7 @@ def __init__( deps_type=deps_type, output_type=[AgentOutput, DeferredToolRequests], tools=[ - Tool(webapp_analyzer) + Tool(webapp_analyzer), ] ) @@ -55,4 +55,4 @@ async def run( usage=usage, usage_limits=usage_limits, deferred_tool_results=deferred_tool_results - ) \ No newline at end of file + ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/recon_threatmodel_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/recon_threatmodel_agent.py index 11d7ce4..c4c77c0 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/recon_threatmodel_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/recon_threatmodel_agent.py @@ -45,7 +45,9 @@ async def run( message_history, usage: RunUsage | None, usage_limits: UsageLimits | None, - deferred_tool_results: DeferredToolResults | None = None + deferred_tool_results: DeferredToolResults | None = None, + *args, + **kwargs ): return await super().run( diff --git a/deadend_cli/deadend_agent/src/deadend_agent/config/settings.py b/deadend_cli/deadend_agent/src/deadend_agent/config/settings.py index 800ba94..debd79f 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/config/settings.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/config/settings.py @@ -341,7 +341,7 @@ def all_model_providers(cls) -> ProvidersList: return cls.providers @classmethod - def get_local_agent_id(cls) -> str: + def get_local_agent_id(cls) -> uuid.UUID: """Return a stable local agent ID, generating one on first use. The ID is persisted in ``config.json`` under the ``local_agent_id`` key @@ -351,10 +351,10 @@ def get_local_agent_id(cls) -> str: config_file = load_config_json() existing = config_file.get("local_agent_id") if existing: - return existing + return uuid.UUID(str(existing)) - new_id = str(uuid.uuid4()) - config_file["local_agent_id"] = new_id + new_id = uuid.uuid4() + config_file["local_agent_id"] = str(new_id) try: _CACHE_TOML_PATH.parent.mkdir(parents=True, exist_ok=True) with open(str(_CACHE_TOML_PATH), "w", encoding="utf-8") as f: diff --git a/deadend_cli/deadend_agent/src/deadend_agent/context/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/context/__init__.py index 66c7316..f62f65f 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/context/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/context/__init__.py @@ -1,7 +1,7 @@ # Copyright (C) 2025 Yassine Bargach # Licensed under the GNU Affero General Public License v3 # See LICENSE file for full license information. -# from .memory import MemoryHandler +from .memory import MemoryHandler from .context_engine import ( ContextEngine, StructuredContext, @@ -10,7 +10,7 @@ ) __all__ = [ - # "MemoryHandler", + "MemoryHandler", "ContextEngine", "StructuredContext", "DiscoveredFact", diff --git a/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py b/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py index cded443..b54733d 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py @@ -21,7 +21,6 @@ from deadend_agent.utils.functions import num_tokens_from_string if TYPE_CHECKING: - from deadend_agent.agents import RouterOutput from deadend_agent.agents.reporter import ReporterAgent @@ -1136,36 +1135,36 @@ async def maybe_summarize_context( self.workflow_context = result.output return token_count - def add_next_agent(self, router_output: "RouterOutput") -> None: - """Add router output information and set the next agent. +# def add_next_agent(self, router_output: "RouterOutput") -> None: +# """Add router output information and set the next agent. - Args: - router_output (RouterOutput): The output from the router agent - containing the next agent name and - routing information. +# Args: +# router_output (RouterOutput): The output from the router agent +# containing the next agent name and +# routing information. - Updates the next_agent attribute and adds the router output - to the workflow context. Also saves to text file. - """ - self.next_agent = router_output.next_agent_name - self.workflow_context += f"""\n -[router agent] -{str(router_output)} -""" - self._append_to_context_file("[ai agent]", f"Router agent: {str(router_output)}") - def add_not_found_agent(self, agent_name: str) -> None: - """Add information about a not found agent to the workflow context. +# Updates the next_agent attribute and adds the router output +# to the workflow context. Also saves to text file. +# """ +# self.next_agent = router_output.next_agent_name +# self.workflow_context += f"""\n +# [router agent] +# {str(router_output)} +# """ +# self._append_to_context_file("[ai agent]", f"Router agent: {str(router_output)}") +# def add_not_found_agent(self, agent_name: str) -> None: +# """Add information about a not found agent to the workflow context. - Args: - agent_name (str): The name of the agent that was not found. +# Args: +# agent_name (str): The name of the agent that was not found. - Adds a message to the workflow context indicating that the - specified agent was not found. Also saves to text file. - """ - self.workflow_context += f""" -[agent not found {agent_name}]\n -""" - self._append_to_context_file("[ai agent]", f"Not found agent name: {agent_name}") +# Adds a message to the workflow context indicating that the +# specified agent was not found. Also saves to text file. +# """ +# self.workflow_context += f""" +# [agent not found {agent_name}]\n +# """ +# self._append_to_context_file("[ai agent]", f"Not found agent name: {agent_name}") def add_agent_response( self, response: str, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/context/memory.py b/deadend_cli/deadend_agent/src/deadend_agent/context/memory.py index dbbfbe9..1ac6fe8 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/context/memory.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/context/memory.py @@ -1,111 +1,66 @@ -# # Copyright (C) 2025 Yassine Bargach -# # Licensed under the GNU Affero General Public License v3 -# # See LICENSE file for full license information. - -# """Memory management system for AI agent conversations and context. - -# This module provides memory management functionality for storing and -# retrieving conversation history, context, and session data for AI agents -# in the security research framework. -# """ -# import json -# from pathlib import Path -# from attr import filters -# from mem0 import MemoryClient - - -# class MemoryHandler: -# """Manages memory storage and retrieval for AI agent conversations and context. - -# This class handles both persistent storage (via mem0 MemoryClient) and local -# file-based caching of agent conversations, tool results, and session data. - -# Attributes: -# memory: MemoryClient instance for persistent memory storage. -# session: Session identifier for the current execution session. -# messages: List of messages in the current conversation. -# target: Target identifier for memory operations. -# base_cache: Path to the base cache directory for this session. -# """ -# memory: MemoryClient - -# def __init__(self, session, target): -# """Initialize the MemoryHandler for a given session and target. - -# Args: -# session: Unique session identifier for this execution session. -# target: Target identifier used for grouping related memories. -# """ -# self.session = session -# self.messages = [] -# self.target = target -# self.base_cache = Path.home() / ".cache" / "deadend" / "memory" / "sessions" / str(self.session) -# self.base_cache.mkdir(parents=True, exist_ok=True) - -# def setup_memory_for_session(self, api_key: str): -# """Initialize the MemoryClient with the provided API key. - -# Args: -# api_key: API key for authenticating with the mem0 memory service. -# """ -# self.memory = MemoryClient(api_key=api_key) - -# def add_agent_conversations(self, messages): -# """Add agent conversation messages to persistent memory storage. - -# Args: -# messages: List of message objects or dictionaries to store in memory. -# Messages are associated with the current target ID. -# """ -# self.memory.add(messages=messages, target_id=self.target) - -# def save_tool_results(self, tool_name: str, **kwargs): -# """Save tool execution results to a local JSONL cache file. - -# Creates a directory structure for the tool and appends results to a -# JSONL file for later retrieval and analysis. - -# Args: -# tool_name: Name of the tool whose results are being saved. -# **kwargs: Additional keyword arguments representing tool-specific -# result data to be saved. -# """ -# tool_cache = self.base_cache / f"{tool_name}" -# tool_cache.mkdir(parents=True, exist_ok=True) -# log_path = tool_cache / f"{tool_name}.jsonl" -# record = {} -# for key, value in kwargs.items(): -# record[key] = value - -# with open(log_path, 'a', encoding="utf-8") as f: -# f.write(json.dumps(record, ensure_ascii=False) + "\n") - -# def add_agent_result_to_memory(self, agent_name: str, **kwargs): -# """Add agent execution results to both persistent memory and local cache. - -# Stores agent results in two locations: -# 1. Persistent memory via MemoryClient (for retrieval by AI) -# 2. Local JSONL file cache (for debugging and analysis) - -# Args: -# agent_name: Name of the agent whose results are being stored. -# **kwargs: Additional keyword arguments representing agent-specific -# result data to be stored. -# """ -# records = {} -# records["agent_name"] = agent_name -# for key, value in kwargs.items(): -# records[key] = value -# self.memory.add(messages=records, target_id=self.target) -# agent_cache = self.base_cache / f"{agent_name}" -# agent_cache.mkdir(parents=True, exist_ok=True) -# log_path = agent_cache / f"{agent_name}.jsonl" -# with open(log_path, 'a', encoding="utf-8") as f: -# f.write(json.dumps(records, ensure_ascii=False) + "\n") - -# def search(self, query: str): -# """Search in memory -# Searches query in memory corresponding to the target in place. -# """ -# result = self.memory.search(query=query, filters={"target_id": self.target}) -# return result +# Copyright (C) 2025 Yassine Bargach +# Licensed under the GNU Affero General Public License v3 +# See LICENSE file for full license information. + +"""High-level memory access built on top of RLM file memory.""" +from __future__ import annotations + +from pathlib import Path +from uuid import UUID + +from deadend_agent.rlm.compat import ( + RLMSandboxCompatibilityReport, + assess_python_sandbox_compatibility, +) +from deadend_agent.rlm.memory import RLMFileMemory, MemoryFileMetadata + + +class MemoryHandler: + """Session-backed memory handler for RLM-style external memory access.""" + + def __init__(self, memory_root: str | Path) -> None: + self.memory_root = Path(memory_root).expanduser().resolve() + self.memory_root.mkdir(parents=True, exist_ok=True) + self._memory = RLMFileMemory(root=self.memory_root) + + @classmethod + def for_session( + cls, + session_key: str | None = None, + session_id: str | UUID | None = None, + base_dir: str | Path | None = None, + ) -> "MemoryHandler": + """Create a handler for a session-backed memory directory.""" + if session_key: + leaf = session_key + elif session_id is not None: + leaf = str(session_id) + else: + raise ValueError("session_key or session_id must be provided") + + root = Path(base_dir) if base_dir else Path.home() / ".cache" / "deadend" / "memory" / "sessions" + return cls(root / leaf) + + def refresh(self) -> None: + """Reload file discovery from disk.""" + self._memory = RLMFileMemory(root=self.memory_root) + + def list_files(self, file_type: str | None = None) -> list[MemoryFileMetadata]: + """Return memory file metadata.""" + return self._memory.list_files(file_type=file_type) + + def describe_memory(self) -> str: + """Return a prompt-friendly navigation summary.""" + return self._memory.build_navigation_context() + + def describe_context(self) -> dict: + """Return prompt metadata for the indexed memory.""" + return self._memory.describe_context() + + def get_rlm_memory(self) -> RLMFileMemory: + """Expose the underlying RLM file memory implementation.""" + return self._memory + + def sandbox_compatibility(self) -> RLMSandboxCompatibilityReport: + """Return compatibility information for the current sandbox backend.""" + return assess_python_sandbox_compatibility() diff --git a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/__init__.py index 14f6245..553c24a 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/__init__.py @@ -61,12 +61,15 @@ class InvalidRequestError(LLMError): # Import main classes from .core_agent import CoreAgent, AgentResult, TokenUsageInfo +from .rlm_runner import RLMRunResult, SandboxedRLMRunner from .session_metrics import SessionMetrics, get_session_metrics, TokenUsage __all__ = [ "CoreAgent", "AgentResult", "TokenUsageInfo", + "SandboxedRLMRunner", + "RLMRunResult", "SessionMetrics", "TokenUsage", "get_session_metrics", diff --git a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/rlm_runner.py b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/rlm_runner.py new file mode 100644 index 0000000..140617e --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/rlm_runner.py @@ -0,0 +1,716 @@ +# Copyright (C) 2025 Yassine Bargach +# Licensed under the GNU Affero General Public License v3 +# See LICENSE file for full license information. + +"""Sandbox-backed Recursive Language Model runner. + +This module implements a practical RLM scaffold for DeadEnd: +- the root LM runs on the host +- the root LM emits Python code blocks +- those code blocks are executed in the existing Python sandbox backend +- sub-LLM calls are orchestrated by the host and exposed to sandboxed code + through a queued ``llm_query(...)`` interface + +The current sandbox backend does not support host callbacks or an in-process +interactive REPL, so this runner emulates a persistent REPL by reusing a +workspace directory and snapshotting picklable globals between iterations. +""" +from __future__ import annotations + +import json +import pickle +import re +import shutil +import textwrap +import uuid +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential + +from deadend_agent.core_agent import ( + AuthenticationError, + ConnectionError, + InvalidRequestError, + LLMError, + ModelNotFoundError, + QuotaExceededError, + RateLimitError, +) +from deadend_agent.rlm.memory import RLMFileMemory, SUPPORTED_EXTENSIONS +from deadend_agent.tools.python_interpreter.python_interpreter import PythonInterpreter + +try: + from litellm import acompletion + from litellm.exceptions import ( + APIConnectionError as LiteLLMConnectionError, + ContentPolicyViolationError, + RateLimitError as LiteLLMRateLimitError, + ServiceUnavailableError, + Timeout as LiteLLMTimeout, + ) + LITELLM_AVAILABLE = True +except ImportError: + acompletion = None + + class LiteLLMConnectionError(Exception): + pass + + class ContentPolicyViolationError(Exception): + pass + + class LiteLLMRateLimitError(Exception): + pass + + class ServiceUnavailableError(Exception): + pass + + class LiteLLMTimeout(Exception): + pass + + LITELLM_AVAILABLE = False + + +PYTHON_BLOCK_RE = re.compile(r"```(?:python|repl)\s*\n(.*?)```", re.DOTALL | re.IGNORECASE) +FINAL_TEXT_RE = re.compile(r"FINAL\((.*)\)", re.DOTALL) +FINAL_VAR_RE = re.compile(r"FINAL_VAR\(\s*([A-Za-z_][A-Za-z0-9_]*)\s*\)") + + +@dataclass +class SubcallRequest: + """A queued sub-LLM request emitted from sandbox code.""" + + request_id: str + prompt: str + content: str = "" + model: str | None = None + + +@dataclass +class SandboxExecutionResult: + """Result of executing one root-LM cell in the sandbox.""" + + stdout: str = "" + observations: list[str] = field(default_factory=list) + state_keys: list[str] = field(default_factory=list) + pending_subcalls: list[SubcallRequest] = field(default_factory=list) + final_answer: str | None = None + final_var: str | None = None + errors: list[str] = field(default_factory=list) + + +@dataclass +class RLMRunResult: + """Final output and metrics for an RLM run.""" + + answer: str + root_iterations: int + root_requests: int + subcall_requests: int + workspace_dir: str + + +class SandboxedRLMRunner: + """Run an RLM loop using host-side LLM calls and sandbox-side Python cells.""" + + def __init__( + self, + root_model: str, + sub_model: str | None = None, + *, + api_key: str | None = None, + api_base: str | None = None, + sub_api_key: str | None = None, + sub_api_base: str | None = None, + session_id: str | None = None, + workspace_root: str | Path | None = None, + ) -> None: + self.root_model = root_model + self.sub_model = sub_model or root_model + self.api_key = api_key + self.api_base = api_base + self.sub_api_key = sub_api_key or api_key + self.sub_api_base = sub_api_base or api_base + self.session_id = session_id or str(uuid.uuid4()) + base = Path(workspace_root) if workspace_root else Path.home() / ".cache" / "deadend" / "python" / "rlm" + self.workspace_dir = base / self.session_id + self.memory_dir = self.workspace_dir / "memory" + self.scripts_dir = self.workspace_dir / "scripts" + self.artifacts_dir = self.workspace_dir / "artifacts" + self.root_request_count = 0 + self.subcall_request_count = 0 + + async def run( + self, + query: str, + *, + memory_root: str | Path | None = None, + context: Any | None = None, + max_iterations: int = 20, + ) -> RLMRunResult: + """Run the root/sub-call RLM loop. + + Args: + query: User query/task for the RLM. + memory_root: Optional existing memory directory to expose. + context: Optional raw context to materialize into memory files. + max_iterations: Safety limit for root iterations. + """ + self._prepare_workspace(memory_root=memory_root, context=context) + memory = RLMFileMemory(self.memory_dir) + messages: list[dict[str, str]] = [ + {"role": "system", "content": self._build_system_prompt(memory)}, + {"role": "user", "content": query}, + ] + + interpreter = PythonInterpreter(session_id=self.session_id, directory=str(self.workspace_dir)) + await interpreter.initialize() + + try: + for iteration in range(1, max_iterations + 1): + assistant_text = await self._call_model( + model=self.root_model, + messages=messages, + api_key=self.api_key, + api_base=self.api_base, + ) + self.root_request_count += 1 + + direct_final = self._parse_direct_final(assistant_text) + if direct_final is not None: + return RLMRunResult( + answer=direct_final, + root_iterations=iteration, + root_requests=self.root_request_count, + subcall_requests=self.subcall_request_count, + workspace_dir=str(self.workspace_dir), + ) + + code_blocks = self._extract_code_blocks(assistant_text) + messages.append({"role": "assistant", "content": assistant_text}) + + if not code_blocks: + messages.append( + { + "role": "user", + "content": ( + "No Python code block was produced. Continue by sending a ```python``` block " + "for the sandbox, or return FINAL(...)." + ), + } + ) + continue + + execution_summaries: list[str] = [] + for cell_index, code in enumerate(code_blocks, start=1): + execution = await self._execute_cell( + interpreter=interpreter, + code=code, + iteration=iteration, + cell_index=cell_index, + ) + await self._resolve_subcalls(execution.pending_subcalls) + execution_summaries.append(self._format_execution_feedback(execution)) + + if execution.final_answer is not None: + return RLMRunResult( + answer=execution.final_answer, + root_iterations=iteration, + root_requests=self.root_request_count, + subcall_requests=self.subcall_request_count, + workspace_dir=str(self.workspace_dir), + ) + + messages.append({"role": "user", "content": "\n\n".join(execution_summaries)}) + + finally: + await interpreter.shutdown() + + raise RuntimeError(f"RLM did not terminate after {max_iterations} root iterations") + + def _prepare_workspace(self, *, memory_root: str | Path | None, context: Any | None) -> None: + self.workspace_dir.mkdir(parents=True, exist_ok=True) + self.memory_dir.mkdir(parents=True, exist_ok=True) + self.scripts_dir.mkdir(parents=True, exist_ok=True) + self.artifacts_dir.mkdir(parents=True, exist_ok=True) + + if memory_root: + source = Path(memory_root).expanduser().resolve() + if not source.exists(): + raise FileNotFoundError(f"Memory root does not exist: {source}") + self._copy_supported_files(source, self.memory_dir) + + if context is not None: + self._materialize_context(context) + + def _copy_supported_files(self, source: Path, destination: Path) -> None: + for path in source.rglob("*"): + if not path.is_file(): + continue + if path.suffix.lower() not in SUPPORTED_EXTENSIONS: + continue + relative = path.relative_to(source) + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, target) + + def _materialize_context(self, context: Any) -> None: + if isinstance(context, str): + (self.memory_dir / "context.txt").write_text(context, encoding="utf-8") + return + if isinstance(context, (dict, list)): + (self.memory_dir / "context.json").write_text( + json.dumps(context, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + return + (self.memory_dir / "context.txt").write_text(str(context), encoding="utf-8") + + def _build_system_prompt(self, memory: RLMFileMemory) -> str: + manifest = memory.build_navigation_context() + return textwrap.dedent( + f"""\ + You are operating as a Recursive Language Model with a sandboxed Python workspace. + + Your job is to answer the user query by inspecting external memory through Python code. + Do not paste long memory contents into the model output. Use the sandbox helpers instead. + + Memory manifest: + {manifest} + + The sandbox runs Python scripts in a persistent workspace. Picklable top-level variables + survive across turns. + + Available helpers inside the sandbox: + - `list_files()` + - `read_file(path)` + - `read_chars(path, start=0, end=None)` + - `read_lines(path, start=1, end=None)` + - `grep_memory(pattern, path=None, max_results=20)` + - `observe(*args)` to emit information back to you + - `subcall_results` containing completed sub-LLM answers by request id + - `llm_query(prompt, content="", model=None)` + - `FINAL(text)` + - `FINAL_VAR(var_name)` + + Important constraints: + - `llm_query(...)` is host-mediated. If the exact request was already completed, it returns the result. + - Otherwise it queues the request and returns a token like `__RLM_PENDING__:...`. + - When a request is pending, stop that line of reasoning and continue on the next turn after inspecting the returned observations. + - Always send Python in fenced blocks: ```python ... ``` + - Return `FINAL(...)` directly only if you do not need the sandbox anymore. + """ + ).strip() + + def _extract_code_blocks(self, text: str) -> list[str]: + return [match.group(1).strip() for match in PYTHON_BLOCK_RE.finditer(text) if match.group(1).strip()] + + def _parse_direct_final(self, text: str) -> str | None: + match = FINAL_VAR_RE.search(text) + if match: + return self._resolve_state_var(match.group(1)) + + match = FINAL_TEXT_RE.search(text) + if not match: + return None + content = match.group(1).strip() + if (content.startswith('"') and content.endswith('"')) or (content.startswith("'") and content.endswith("'")): + return content[1:-1] + return content + + def _resolve_state_var(self, var_name: str) -> str | None: + state_path = self.artifacts_dir / "globals.pkl" + if not state_path.exists(): + return None + try: + with open(state_path, "rb") as file_obj: + state = pickle.load(file_obj) + except Exception: + return None + if not isinstance(state, dict) or var_name not in state: + return None + value = state[var_name] + return value if isinstance(value, str) else repr(value) + + async def _execute_cell( + self, + *, + interpreter: PythonInterpreter, + code: str, + iteration: int, + cell_index: int, + ) -> SandboxExecutionResult: + script_name = f"rlm_iter_{iteration:02d}_cell_{cell_index:02d}.py" + script_path = self.scripts_dir / script_name + script_path.write_text(self._build_script(code), encoding="utf-8") + + raw_result = await interpreter.run_file(str(Path("scripts") / script_name)) + observation = self._read_json_artifact("observation.json", default={}) + pending = self._read_json_artifact("pending_subcalls.json", default=[]) + final_payload = self._read_json_artifact("final.json", default=None) + + final_answer: str | None = None + final_var: str | None = None + if isinstance(final_payload, dict): + final_answer = final_payload.get("value") + final_var = final_payload.get("var_name") + + return SandboxExecutionResult( + stdout=str(raw_result), + observations=list(observation.get("observations", [])), + state_keys=list(observation.get("state_keys", [])), + pending_subcalls=[ + SubcallRequest( + request_id=item["request_id"], + prompt=item["prompt"], + content=item.get("content", ""), + model=item.get("model"), + ) + for item in pending + ], + final_answer=final_answer, + final_var=final_var, + errors=list(observation.get("errors", [])), + ) + + def _build_script(self, user_code: str) -> str: + workspace = str(self.workspace_dir) + memory_dir = str(self.memory_dir) + artifacts_dir = str(self.artifacts_dir) + indented_code = textwrap.indent(user_code.rstrip() + "\n", " ").rstrip() + template = textwrap.dedent( + f"""\ + import hashlib + import json + import os + import pickle + import re + import traceback + from pathlib import Path + + WORKSPACE_DIR = Path(r"{workspace}") + MEMORY_DIR = Path(r"{memory_dir}") + ARTIFACTS_DIR = Path(r"{artifacts_dir}") + STATE_PATH = ARTIFACTS_DIR / "globals.pkl" + SUBCALL_RESULTS_PATH = ARTIFACTS_DIR / "subcall_results.json" + PENDING_SUBCALLS_PATH = ARTIFACTS_DIR / "pending_subcalls.json" + OBSERVATION_PATH = ARTIFACTS_DIR / "observation.json" + FINAL_PATH = ARTIFACTS_DIR / "final.json" + + for _directory in (WORKSPACE_DIR, MEMORY_DIR, ARTIFACTS_DIR): + _directory.mkdir(parents=True, exist_ok=True) + + _BOOTSTRAP_NAMES = {{ + "hashlib", "json", "os", "pickle", "re", "traceback", "Path", + "WORKSPACE_DIR", "MEMORY_DIR", "ARTIFACTS_DIR", "STATE_PATH", + "SUBCALL_RESULTS_PATH", "PENDING_SUBCALLS_PATH", "OBSERVATION_PATH", + "FINAL_PATH", "_BOOTSTRAP_NAMES", "_load_json", "_write_json", + "_runtime", "_stable_subcall_id", "_resolve_memory_path", "_save_state", + "_load_state", "subcall_results", "observe", "llm_query", "FINAL", + "FINAL_VAR", "list_files", "read_file", "read_chars", "read_lines", + "grep_memory", "_supported_suffixes" + }} + + _supported_suffixes = {sorted(SUPPORTED_EXTENSIONS)!r} + + def _load_json(path, default): + if not path.exists(): + return default + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return default + + def _write_json(path, payload): + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + + subcall_results = _load_json(SUBCALL_RESULTS_PATH, {{}}) + _runtime = {{ + "observations": [], + "pending_subcalls": [], + "errors": [], + "final": None, + }} + + def _stable_subcall_id(prompt, content="", model=None): + payload = json.dumps( + {{"prompt": prompt, "content": content, "model": model or ""}}, + sort_keys=True, + ensure_ascii=False, + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:24] + + def _resolve_memory_path(path): + candidate = (MEMORY_DIR / path).resolve() + if not str(candidate).startswith(str(MEMORY_DIR.resolve())): + raise ValueError(f"Path escapes memory directory: {{path}}") + if not candidate.exists(): + raise FileNotFoundError(f"Memory file not found: {{path}}") + return candidate + + def list_files(): + output = [] + for file_path in sorted(MEMORY_DIR.rglob("*")): + if file_path.is_file() and file_path.suffix.lower() in _supported_suffixes: + rel = file_path.relative_to(MEMORY_DIR) + output.append(str(rel)) + return output + + def read_file(path): + return _resolve_memory_path(path).read_text(encoding="utf-8") + + def read_chars(path, start=0, end=None): + text = read_file(path) + return text[start:end] + + def read_lines(path, start=1, end=None): + lines = read_file(path).splitlines() + start_index = max(0, int(start) - 1) + end_index = len(lines) if end is None else max(start_index, int(end)) + return "\\n".join(lines[start_index:end_index]) + + def grep_memory(pattern, path=None, max_results=20): + compiled = re.compile(pattern, re.MULTILINE) + targets = [path] if path else list_files() + matches = [] + for target in targets: + for line_number, line in enumerate(read_file(target).splitlines(), start=1): + match = compiled.search(line) + if not match: + continue + matches.append({{ + "path": target, + "line_number": line_number, + "match": match.group(0), + "context": line[:240], + }}) + if len(matches) >= max_results: + return matches + return matches + + def observe(*args): + text = " ".join(str(arg) for arg in args) + _runtime["observations"].append(text) + print(text) + return text + + def llm_query(prompt, content="", model=None): + request_id = _stable_subcall_id(prompt, content=content, model=model) + if request_id in subcall_results: + return subcall_results[request_id]["result"] + _runtime["pending_subcalls"].append({{ + "request_id": request_id, + "prompt": prompt, + "content": content, + "model": model, + }}) + token = f"__RLM_PENDING__:{{request_id}}" + observe("queued_subcall", request_id) + return token + + def FINAL(text): + payload = {{"type": "text", "value": str(text)}} + _runtime["final"] = payload + _write_json(FINAL_PATH, payload) + return text + + def FINAL_VAR(var_name): + if var_name not in globals(): + raise KeyError(f"Variable not found for FINAL_VAR: {{var_name}}") + value = globals()[var_name] + payload = {{ + "type": "var", + "var_name": var_name, + "value": value if isinstance(value, str) else repr(value), + }} + _runtime["final"] = payload + _write_json(FINAL_PATH, payload) + return value + + def _load_state(): + if not STATE_PATH.exists(): + return + try: + with open(STATE_PATH, "rb") as file_obj: + state = pickle.load(file_obj) + if isinstance(state, dict): + for key, value in state.items(): + globals()[key] = value + except Exception as exc: + _runtime["errors"].append(f"state_load_error: {{exc!r}}") + + def _save_state(): + state = {{}} + for key, value in list(globals().items()): + if key.startswith("_") or key in _BOOTSTRAP_NAMES: + continue + if callable(value): + continue + if getattr(value, "__class__", None).__name__ == "module": + continue + try: + pickle.dumps(value) + except Exception: + continue + state[key] = value + with open(STATE_PATH, "wb") as file_obj: + pickle.dump(state, file_obj) + return sorted(state.keys()) + + _load_state() + FINAL_PATH.unlink(missing_ok=True) + PENDING_SUBCALLS_PATH.unlink(missing_ok=True) + + try: + __USER_CODE__ + except Exception: + _runtime["errors"].append(traceback.format_exc()) + finally: + state_keys = _save_state() + _write_json(PENDING_SUBCALLS_PATH, _runtime["pending_subcalls"]) + _write_json( + OBSERVATION_PATH, + {{ + "observations": _runtime["observations"], + "errors": _runtime["errors"], + "state_keys": state_keys, + "final": _runtime["final"], + }}, + ) + """ + ).lstrip() + return template.replace("__USER_CODE__", indented_code) + + async def _resolve_subcalls(self, requests: list[SubcallRequest]) -> None: + if not requests: + return + + results = self._read_json_artifact("subcall_results.json", default={}) + for request in requests: + if request.request_id in results: + continue + result = await self._call_model( + model=request.model or self.sub_model, + messages=self._build_subcall_messages(request), + api_key=self.sub_api_key, + api_base=self.sub_api_base, + ) + self.subcall_request_count += 1 + results[request.request_id] = { + "prompt": request.prompt, + "content": request.content, + "model": request.model or self.sub_model, + "result": result, + } + + self._write_json_artifact("subcall_results.json", results) + + def _build_subcall_messages(self, request: SubcallRequest) -> list[dict[str, str]]: + content = request.prompt + if request.content: + content += f"\n\nContext:\n{request.content}" + return [{"role": "user", "content": content}] + + def _format_execution_feedback(self, execution: SandboxExecutionResult) -> str: + parts = [] + if execution.observations: + parts.append("Sandbox observations:\n" + "\n".join(f"- {item}" for item in execution.observations[-20:])) + if execution.pending_subcalls: + parts.append( + "Queued subcalls:\n" + + "\n".join(f"- {item.request_id}" for item in execution.pending_subcalls) + ) + sub_results = self._read_json_artifact("subcall_results.json", default={}) + completed = [ + f"- {item.request_id}: {sub_results[item.request_id]['result'][:500]}" + for item in execution.pending_subcalls + if item.request_id in sub_results + ] + if completed: + parts.append("Completed subcall results:\n" + "\n".join(completed)) + if execution.errors: + parts.append("Sandbox errors:\n" + "\n".join(execution.errors)) + if execution.state_keys: + parts.append("Persisted state keys:\n" + ", ".join(execution.state_keys)) + stdout = execution.stdout.strip() + if stdout: + parts.append("Sandbox backend response:\n" + stdout[:2000]) + return "\n\n".join(parts) if parts else "Sandbox cell executed with no observations." + + def _read_json_artifact(self, filename: str, default: Any) -> Any: + path = self.artifacts_dir / filename + if not path.exists(): + return default + try: + return json.loads(path.read_text(encoding="utf-8")) + except Exception: + return default + + def _write_json_artifact(self, filename: str, payload: Any) -> None: + path = self.artifacts_dir / filename + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + + @retry( + stop=stop_after_attempt(5), + wait=wait_exponential(multiplier=2, min=2, max=30), + retry=retry_if_exception_type( + ( + LiteLLMRateLimitError, + ServiceUnavailableError, + LiteLLMTimeout, + LiteLLMConnectionError, + ) + ), + reraise=True, + ) + async def _call_model( + self, + *, + model: str, + messages: list[dict[str, str]], + api_key: str | None, + api_base: str | None, + ) -> str: + if not LITELLM_AVAILABLE or acompletion is None: + raise RuntimeError("litellm is required to run SandboxedRLMRunner model calls") + + kwargs: dict[str, Any] = { + "model": model, + "messages": messages, + } + if api_base: + kwargs["api_base"] = api_base + if api_key: + kwargs["api_key"] = api_key + elif api_base and model.startswith("openai/"): + kwargs["api_key"] = "sk-dummy-key-for-local-model" + + try: + response = await acompletion(**kwargs) + except ContentPolicyViolationError as exc: + raise InvalidRequestError( + "Request blocked by provider content policy.", + original_error=exc, + ) from exc + except Exception as exc: + self._raise_llm_error(model=model, error=exc) + + content = response.choices[0].message.content + return content or "" + + def _raise_llm_error(self, *, model: str, error: Exception) -> None: + error_str = str(error).lower() + if "insufficient_quota" in error_str or "exceeded your current quota" in error_str: + raise QuotaExceededError(f"API quota exceeded for {model}: {error}", original_error=error) from error + if "rate_limit" in error_str or "rate limit" in error_str or "429" in error_str: + raise RateLimitError(f"Rate limit exceeded for {model}: {error}", original_error=error) from error + if "auth" in error_str or "api_key" in error_str or "401" in error_str: + raise AuthenticationError(f"API authentication failed for {model}: {error}", original_error=error) from error + if "model" in error_str and ("not found" in error_str or "404" in error_str or "does not exist" in error_str): + raise ModelNotFoundError(f"Model '{model}' not found: {error}", original_error=error) from error + if "connection" in error_str or "connect" in error_str or "timeout" in error_str or "unreachable" in error_str: + raise ConnectionError(f"Failed to connect for {model}: {error}", original_error=error) from error + if "bad request" in error_str or "invalid" in error_str or "400" in error_str: + raise InvalidRequestError(f"Invalid request for {model}: {error}", original_error=error) from error + raise LLMError(f"LLM request failed for {model}: {error}", original_error=error) from error diff --git a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py index 9132d8b..da7cdc3 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py @@ -1,12 +1,13 @@ """Main DeadEnd agent orchestration module.""" import re +from pathlib import Path from typing import Any, Awaitable, Callable, Dict, Generator from uuid import UUID from deadend_agent.logging import logger from pydantic_ai import RunUsage, UsageLimits -from deadend_agent.config.settings import ModelSpec +from deadend_agent.config.settings import Config, ModelSpec from deadend_agent.models.registry import EmbedderClient from deadend_agent.embedders.code_indexer import SourceCodeIndexer from deadend_agent.context import ContextEngine @@ -14,16 +15,19 @@ from deadend_agent.sandbox.sandbox import Sandbox from deadend_agent.agents.reporter import ReporterAgent from deadend_agent.agents.architecture import ADaPTAgent +from deadend_agent.agents.generic_agents.memory_agent import MemoryAgent from deadend_agent.agents.components.executor import AgentExecutor, ResultEvent from deadend_agent.agents.components.planner import Planner, TaskNode from deadend_agent.agents.components.validator import Validator from deadend_agent.utils.structures import ( + MemoryWorkspaceDeps, RequesterDeps, ShellDeps, ShellRunner, WebappreconDeps ) from deadend_agent.tools.browser_automation.http_parser import extract_host_port +from deadend_agent.tools.avfs import avfs from .agents.recon_threatmodel_agent import ReconThreatModelAgent from .agents.exploit_web_agent import PlannerExploitAgent # from deadend_eval.metrics imporst save_traces @@ -34,6 +38,7 @@ class DeadEndAgent: """Main orchestrator for the DeadEnd security research framework.""" + agent_id: UUID session_id: UUID embedding_session_id: UUID model: ModelSpec @@ -55,6 +60,7 @@ class DeadEndAgent: requester_deps: RequesterDeps | None = None webapprecon_deps: WebappreconDeps | None = None challenge_name: str | None = None + local_agent_id: UUID def __init__( @@ -65,7 +71,10 @@ def __init__( max_depth: int = 3, validation_type: str | None = None, validation_format: str | None = None, - embedding_session_id: UUID | None = None + embedding_session_id: UUID | None = None, + workspace_root: str | None = None, + agents_storage_root: str | None = None, + local_agent_id: UUID | None = None, ): self.session_id = session_id self.embedding_session_id = embedding_session_id or session_id @@ -78,6 +87,19 @@ def __init__( validation_format=validation_format ) self.context = ContextEngine(model=self.model, session_id=session_id) + self.workspace_root: str | None = None + self.local_agent_id = local_agent_id or Config.get_local_agent_id() + self.agent_id = self.local_agent_id + self.agents_storage_root = agents_storage_root or Config.agents_storage_root + self.memory_workspace_root = self._prepare_memory_workspace() + self.memory_context = "" + avfs.mount( + workspace_root=self.memory_workspace_root, + session_id=str(self.agent_id), + workspace="memory", + ) + if workspace_root is not None: + self.set_workspace_root(workspace_root) ################################################################################ @@ -111,6 +133,64 @@ def set_approval_callback(self, callback): callback: Async function that returns user input for approval """ self.approval_callback = callback + + def set_workspace_root(self, workspace_root: str | None) -> None: + """Configure the AVFS workspace for this agent session.""" + if workspace_root is None: + self.workspace_root = None + avfs.umount(session_id=str(self.agent_id), workspace="workspace") + return + + avfs.mount(workspace_root=workspace_root, session_id=str(self.agent_id), workspace="workspace") + mounted_root = avfs.current_workspace_root(session_id=str(self.agent_id), workspace="workspace") + self.workspace_root = None if mounted_root is None else str(mounted_root) + + def _prepare_memory_workspace(self) -> str: + """Ensure the persistent memory workspace exists for this local agent.""" + memory_root = ( + Path(self.agents_storage_root).expanduser().resolve() + / str(self.local_agent_id) + / str(self.embedding_session_id) + / "memory" + ) + memory_root.mkdir(parents=True, exist_ok=True) + return str(memory_root) + + async def _populate_memory_context(self, task_query: str) -> str: + """Refresh task-specific memory context right before supervisor execution.""" + memory_agent = MemoryAgent( + model=self.model, + deps_type=MemoryWorkspaceDeps, + ) + memory_deps = MemoryWorkspaceDeps( + session_id=str(self.agent_id), + memory_workspace_root=self.memory_workspace_root, + ) + result = await memory_agent.run( + prompt=( + f"Current task:\n{task_query}\n\n" + "Inspect the persistent memory workspace using AVFS tools with workspace=\"memory\". " + "Return only a concise task-relevant memory summary as plain text for the supervisor. " + "If memory is empty or not useful for this task, return a short plain-text statement saying that no relevant persisted memory is available." + ), + deps=memory_deps, + message_history=[], + usage=RunUsage(), + usage_limits=UsageLimits(request_limit=None, tool_calls_limit=None), + deferred_tool_results=None, + ) + + output = getattr(result, "output", None) + self.memory_context = str(output).strip() if output is not None else "" + if hasattr(self, "executor"): + self.executor.set_memory_context(self.memory_context) + if self.shell_deps is not None: + self.shell_deps.memory_context = self.memory_context + if self.requester_deps is not None: + self.requester_deps.memory_context = self.memory_context + if self.webapprecon_deps is not None: + self.webapprecon_deps.memory_context = self.memory_context + return self.memory_context ################################################################################## ################################################################################## @@ -205,15 +285,23 @@ def prepare_dependencies( raise ValueError("target must be provided before initializing dependencies.") - shell_runner = ShellRunner(session=str(self.session_id), sandbox=sandbox) + shell_runner = ShellRunner(session=str(self.agent_id), sandbox=sandbox) - self.shell_deps = ShellDeps(shell_runner=shell_runner) + self.shell_deps = ShellDeps( + shell_runner=shell_runner, + session_id=self.agent_id, + workspace_root=self.workspace_root, + memory_workspace_root=self.memory_workspace_root, + memory_context=self.memory_context, + ) self.requester_deps = RequesterDeps( embedder_client=embedder_client, rag=rag_connector, target=target_host, session_id=self.session_id, - embedding_session_id=self.embedding_session_id + embedding_session_id=self.embedding_session_id, + memory_workspace_root=self.memory_workspace_root, + memory_context=self.memory_context, ) self.webapprecon_deps = WebappreconDeps( embedder_client=embedder_client, @@ -221,14 +309,18 @@ def prepare_dependencies( target=target_host, shell_runner=shell_runner, session_id=self.session_id, - embedding_session_id=self.embedding_session_id + embedding_session_id=self.embedding_session_id, + memory_workspace_root=self.memory_workspace_root, + memory_context=self.memory_context, ) self.executor = AgentExecutor( model=self.model, context=self.context, available_agents=self.available_agents, - session_id=self._target_session_key() + session_id=str(self.agent_id) ) + self.executor.set_auth_session_key(self._target_session_key()) + self.executor.set_memory_context(self.memory_context) self.executor.set_dependencies( requester_deps=self.requester_deps, @@ -700,4 +792,4 @@ async def start_supervisor(self, task: str): message_history="" ) - yield threat_model_data \ No newline at end of file + yield threat_model_data diff --git a/deadend_cli/deadend_agent/src/deadend_agent/rlm/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/rlm/__init__.py new file mode 100644 index 0000000..e702d82 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/rlm/__init__.py @@ -0,0 +1,16 @@ +from .compat import RLMSandboxCompatibilityReport, assess_python_sandbox_compatibility +from .memory import ( + MarkdownSection, + MemoryFileMetadata, + MemorySearchResult, + RLMFileMemory, +) + +__all__ = [ + "MarkdownSection", + "MemoryFileMetadata", + "MemorySearchResult", + "RLMFileMemory", + "RLMSandboxCompatibilityReport", + "assess_python_sandbox_compatibility", +] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/rlm/compat.py b/deadend_cli/deadend_agent/src/deadend_agent/rlm/compat.py new file mode 100644 index 0000000..0d4ae69 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/rlm/compat.py @@ -0,0 +1,52 @@ +# Copyright (C) 2025 Yassine Bargach +# Licensed under the GNU Affero General Public License v3 +# See LICENSE file for full license information. + +"""Compatibility checks for running RLM flows against current execution backends.""" +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class RLMSandboxCompatibilityReport: + """Concrete capability report for the current Python sandbox integration.""" + + backend_name: str + persistent_state: bool + inline_code_execution: bool + host_callback_support: bool + raw_llm_calls_required_inside_sandbox: bool + compatible_for_full_rlm_repl: bool + blockers: list[str] = field(default_factory=list) + recommendation: str = "" + + +def assess_python_sandbox_compatibility() -> RLMSandboxCompatibilityReport: + """Assess whether the current sandbox can host a faithful RLM REPL. + + The answer is currently no. The existing wrapper starts a fresh sandbox per + tool invocation, executes a file, then shuts the sandbox down. There is no + persistent Python state across turns, no inline REPL API, and no mechanism + for code running inside the sandbox to call a host-managed ``llm_query``. + """ + blockers = [ + "The sandbox wrapper only exposes run-file semantics, not incremental REPL execution.", + "Each tool invocation starts a new interpreter process and shuts it down afterwards.", + "There is no host callback channel for llm_query(prompt, content) from sandboxed code.", + "Using raw provider calls inside sandboxed scripts would bypass CoreAgent and duplicate auth/routing logic.", + ] + return RLMSandboxCompatibilityReport( + backend_name="python-sandbox-tool", + persistent_state=False, + inline_code_execution=False, + host_callback_support=False, + raw_llm_calls_required_inside_sandbox=False, + compatible_for_full_rlm_repl=False, + blockers=blockers, + recommendation=( + "Keep RLM orchestration and LLM calls on the host side. Use the sandbox only as a " + "bounded execution backend after adding either persistent REPL support or an explicit " + "host-mediated llm_query callback protocol." + ), + ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/rlm/memory.py b/deadend_cli/deadend_agent/src/deadend_agent/rlm/memory.py new file mode 100644 index 0000000..ffa6a8d --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/rlm/memory.py @@ -0,0 +1,473 @@ +# Copyright (C) 2025 Yassine Bargach +# Licensed under the GNU Affero General Public License v3 +# See LICENSE file for full license information. + +"""RLM-oriented file memory primitives. + +This module implements the first layer needed for RLM-style memory: +long-lived memory remains outside the prompt as files on disk, and the +agent gets structure-aware operations to inspect that memory selectively. + +The current implementation focuses on the memory substrate itself: +- markdown files are indexed by headings +- JSON files are navigated by path +- JSONL files are exposed as arrays of JSON objects +- callers can build prompt metadata without reading the full corpus +""" +from __future__ import annotations + +from dataclasses import dataclass +import json +import re +from pathlib import Path +from typing import Any, Iterable + + +MARKDOWN_EXTENSIONS = {".md", ".markdown"} +JSON_EXTENSIONS = {".json"} +JSONL_EXTENSIONS = {".jsonl"} +TEXT_EXTENSIONS = {".txt", ".log"} +SUPPORTED_EXTENSIONS = MARKDOWN_EXTENSIONS | JSON_EXTENSIONS | JSONL_EXTENSIONS | TEXT_EXTENSIONS + + +@dataclass(frozen=True) +class MemoryFileMetadata: + """Metadata for a memory file.""" + + path: str + absolute_path: str + file_type: str + size_bytes: int + line_count: int + section_count: int = 0 + + +@dataclass(frozen=True) +class MarkdownSection: + """Indexed markdown section.""" + + section_id: str + heading: str + level: int + start_line: int + end_line: int + char_start: int + char_end: int + content: str + + +@dataclass(frozen=True) +class MemorySearchResult: + """Structured grep-like match.""" + + path: str + line_number: int + match: str + context: str + + +class RLMFileMemory: + """Structure-aware external memory for long markdown and JSON corpora.""" + + def __init__(self, root: str | Path, files: Iterable[str | Path] | None = None) -> None: + self.root = Path(root).expanduser().resolve() + self.root.mkdir(parents=True, exist_ok=True) + self._files: dict[str, Path] = {} + + selected_files = files if files is not None else self._discover_files() + for file_path in selected_files: + path = Path(file_path).expanduser().resolve() + if not path.exists() or not path.is_file(): + continue + if path.suffix.lower() not in SUPPORTED_EXTENSIONS: + continue + key = self._normalize_key(path) + self._files[key] = path + + def _discover_files(self) -> list[Path]: + discovered: list[Path] = [] + for path in self.root.rglob("*"): + if path.is_file() and path.suffix.lower() in SUPPORTED_EXTENSIONS: + discovered.append(path.resolve()) + return sorted(discovered) + + def _normalize_key(self, path: Path) -> str: + try: + return str(path.relative_to(self.root)) + except ValueError: + return path.name + + def _resolve_file(self, path: str) -> Path: + candidate = self._files.get(path) + if candidate: + return candidate + + fallback = (self.root / path).resolve() + if fallback.exists() and fallback.is_file(): + return fallback + + raise FileNotFoundError(f"Memory file not found: {path}") + + def _read_text(self, path: str) -> str: + file_path = self._resolve_file(path) + return file_path.read_text(encoding="utf-8") + + def list_files(self, file_type: str | None = None) -> list[MemoryFileMetadata]: + """Return indexed files with lightweight metadata.""" + items: list[MemoryFileMetadata] = [] + for key, path in sorted(self._files.items()): + kind = self._detect_file_type(path) + if file_type and kind != file_type: + continue + text = path.read_text(encoding="utf-8") + section_count = len(self._split_markdown_sections(text)) if kind == "markdown" else 0 + items.append( + MemoryFileMetadata( + path=key, + absolute_path=str(path), + file_type=kind, + size_bytes=len(text.encode("utf-8")), + line_count=text.count("\n") + (1 if text else 0), + section_count=section_count, + ) + ) + return items + + def get_file_metadata(self, path: str) -> MemoryFileMetadata: + """Return metadata for one file.""" + matches = self.list_files() + for metadata in matches: + if metadata.path == path: + return metadata + raise FileNotFoundError(f"Memory file not found: {path}") + + def grep_memory( + self, + pattern: str, + path: str | None = None, + max_results: int = 50, + ) -> list[MemorySearchResult]: + """Search memory files with a regex pattern.""" + compiled = re.compile(pattern, re.MULTILINE) + targets = [path] if path else [metadata.path for metadata in self.list_files()] + results: list[MemorySearchResult] = [] + + for target in targets: + text = self._read_text(target) + for line_number, line in enumerate(text.splitlines(), start=1): + match = compiled.search(line) + if not match: + continue + results.append( + MemorySearchResult( + path=target, + line_number=line_number, + match=match.group(0), + context=line[:240], + ) + ) + if len(results) >= max_results: + return results + return results + + def read_chars(self, path: str, start: int = 0, end: int | None = None) -> str: + """Read a character slice from a file.""" + text = self._read_text(path) + return text[start:end] + + def read_lines(self, path: str, start: int = 1, end: int | None = None) -> str: + """Read a line slice from a file using 1-based indexing.""" + lines = self._read_text(path).splitlines() + start_index = max(0, start - 1) + end_index = len(lines) if end is None else max(start_index, end) + return "\n".join(lines[start_index:end_index]) + + def list_md_sections(self, path: str) -> list[MarkdownSection]: + """Return markdown sections indexed by heading.""" + file_path = self._resolve_file(path) + if file_path.suffix.lower() not in MARKDOWN_EXTENSIONS: + raise ValueError(f"Not a markdown file: {path}") + return self._split_markdown_sections(self._read_text(path)) + + def read_md_section(self, path: str, heading_or_id: str) -> str: + """Read a markdown section by heading text or section id.""" + for section in self.list_md_sections(path): + if section.section_id == heading_or_id or section.heading == heading_or_id: + return section.content + raise KeyError(f"Markdown section not found: {heading_or_id}") + + def read_md_outline(self, path: str) -> list[dict[str, Any]]: + """Return a condensed outline for a markdown file.""" + return [ + { + "section_id": section.section_id, + "heading": section.heading, + "level": section.level, + "start_line": section.start_line, + "end_line": section.end_line, + } + for section in self.list_md_sections(path) + ] + + def search_md_headings(self, query: str) -> list[dict[str, Any]]: + """Search markdown headings across indexed files.""" + query_lower = query.lower() + matches: list[dict[str, Any]] = [] + for metadata in self.list_files(file_type="markdown"): + for section in self.list_md_sections(metadata.path): + if query_lower in section.heading.lower(): + matches.append( + { + "path": metadata.path, + "section_id": section.section_id, + "heading": section.heading, + "level": section.level, + } + ) + return matches + + def json_keys(self, path: str, json_path: str | None = None) -> list[str]: + """List keys available at a JSON path.""" + node = self._resolve_json_node(path, json_path) + if isinstance(node, dict): + return list(node.keys()) + if isinstance(node, list): + return [str(index) for index in range(len(node))] + return [] + + def json_get(self, path: str, json_path: str | None = None) -> Any: + """Read a JSON value at a path.""" + return self._resolve_json_node(path, json_path) + + def json_search( + self, + path: str, + key: str | None = None, + value_contains: str | None = None, + max_results: int = 50, + ) -> list[dict[str, Any]]: + """Search a JSON structure recursively.""" + data = self._load_json_payload(path) + results: list[dict[str, Any]] = [] + needle = value_contains.lower() if value_contains else None + + def walk(node: Any, current_path: str) -> None: + if len(results) >= max_results: + return + if isinstance(node, dict): + for child_key, child_value in node.items(): + child_path = f"{current_path}.{child_key}" if current_path else child_key + key_match = key is None or child_key == key + value_match = ( + needle is None + or (isinstance(child_value, (str, int, float, bool)) and needle in str(child_value).lower()) + ) + if key_match and value_match: + results.append({"path": child_path, "key": child_key, "value": child_value}) + walk(child_value, child_path) + elif isinstance(node, list): + for index, child_value in enumerate(node): + child_path = f"{current_path}[{index}]" if current_path else f"[{index}]" + if needle is not None and isinstance(child_value, (str, int, float, bool)): + if needle in str(child_value).lower(): + results.append({"path": child_path, "key": None, "value": child_value}) + walk(child_value, child_path) + + walk(data, "") + return results[:max_results] + + def json_sample_array(self, path: str, json_path: str, start: int = 0, end: int = 10) -> list[Any]: + """Return a slice from a JSON array.""" + node = self._resolve_json_node(path, json_path) + if not isinstance(node, list): + raise ValueError(f"JSON path does not resolve to an array: {json_path}") + return node[start:end] + + def json_schema(self, path: str, json_path: str | None = None, max_depth: int = 5) -> Any: + """Return a lightweight structural schema for a JSON value.""" + node = self._resolve_json_node(path, json_path) + return self._infer_schema(node, depth=0, max_depth=max_depth) + + def describe_context(self) -> dict[str, Any]: + """Return prompt-friendly metadata for the external memory corpus.""" + files = self.list_files() + return { + "context_type": "RLMFileMemory", + "context_total_length": sum(item.size_bytes for item in files), + "context_lengths": [item.size_bytes for item in files], + "files": [ + { + "path": item.path, + "file_type": item.file_type, + "size_bytes": item.size_bytes, + "line_count": item.line_count, + "section_count": item.section_count, + } + for item in files + ], + } + + def build_navigation_context(self) -> str: + """Return a compact index that an LLM can inspect before reading content.""" + description = self.describe_context() + lines = [ + f"context_type={description['context_type']}", + f"context_total_length={description['context_total_length']}", + "files:", + ] + for item in description["files"]: + lines.append( + f"- {item['path']} [{item['file_type']}] " + f"size={item['size_bytes']} lines={item['line_count']} sections={item['section_count']}" + ) + return "\n".join(lines) + + def _detect_file_type(self, path: Path) -> str: + suffix = path.suffix.lower() + if suffix in MARKDOWN_EXTENSIONS: + return "markdown" + if suffix in JSON_EXTENSIONS: + return "json" + if suffix in JSONL_EXTENSIONS: + return "jsonl" + return "text" + + def _split_markdown_sections(self, text: str) -> list[MarkdownSection]: + if not text: + return [] + + lines = text.splitlines() + sections: list[MarkdownSection] = [] + heading_pattern = re.compile(r"^(#{1,6})\s+(.*)$") + current_heading = "ROOT" + current_level = 0 + current_start_line = 1 + current_lines: list[str] = [] + + def flush(end_line: int) -> None: + if not current_lines and current_heading == "ROOT": + return + content = "\n".join(current_lines).strip() + char_start = len("\n".join(lines[: current_start_line - 1])) + if current_start_line > 1: + char_start += 1 + char_end = char_start + len(content) + sections.append( + MarkdownSection( + section_id=self._slugify_heading(current_heading, len(sections) + 1), + heading=current_heading, + level=current_level, + start_line=current_start_line, + end_line=end_line, + char_start=char_start, + char_end=char_end, + content=content, + ) + ) + + for line_number, line in enumerate(lines, start=1): + match = heading_pattern.match(line) + if match: + flush(line_number - 1) + current_heading = match.group(2).strip() + current_level = len(match.group(1)) + current_start_line = line_number + current_lines = [line] + continue + + current_lines.append(line) + + flush(len(lines)) + return sections + + def _slugify_heading(self, heading: str, index: int) -> str: + slug = re.sub(r"[^a-z0-9]+", "-", heading.lower()).strip("-") + return slug or f"section-{index}" + + def _load_json_payload(self, path: str) -> Any: + file_path = self._resolve_file(path) + text = file_path.read_text(encoding="utf-8") + if file_path.suffix.lower() in JSONL_EXTENSIONS: + return [json.loads(line) for line in text.splitlines() if line.strip()] + return json.loads(text) + + def _resolve_json_node(self, path: str, json_path: str | None) -> Any: + node = self._load_json_payload(path) + if not json_path or json_path in {".", "$"}: + return node + + segments = self._parse_json_path(json_path) + for segment in segments: + if isinstance(segment, int): + if not isinstance(node, list): + raise KeyError(f"Expected list while resolving index {segment} in {json_path}") + node = node[segment] + continue + if not isinstance(node, dict): + raise KeyError(f"Expected object while resolving key '{segment}' in {json_path}") + node = node[segment] + return node + + def _parse_json_path(self, json_path: str) -> list[str | int]: + normalized = json_path.strip() + if normalized.startswith("$"): + normalized = normalized[1:] + normalized = normalized.lstrip(".") + + if not normalized: + return [] + + segments: list[str | int] = [] + token = "" + index_token = "" + in_index = False + + for char in normalized: + if char == "." and not in_index: + if token: + segments.append(token) + token = "" + continue + if char == "[": + if token: + segments.append(token) + token = "" + in_index = True + index_token = "" + continue + if char == "]" and in_index: + segments.append(int(index_token)) + index_token = "" + in_index = False + continue + if in_index: + index_token += char + else: + token += char + + if token: + segments.append(token) + + return segments + + def _infer_schema(self, node: Any, depth: int, max_depth: int) -> Any: + if depth >= max_depth: + return {"type": type(node).__name__} + if isinstance(node, dict): + return { + "type": "object", + "properties": { + key: self._infer_schema(value, depth + 1, max_depth) + for key, value in node.items() + }, + } + if isinstance(node, list): + item_schema = self._infer_schema(node[0], depth + 1, max_depth) if node else {"type": "unknown"} + return { + "type": "array", + "length": len(node), + "items": item_schema, + } + if node is None: + return {"type": "null"} + return {"type": type(node).__name__} diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py index ad7ffd5..9ed6c39 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py @@ -22,6 +22,7 @@ from .grep import grep_session_logs from .webapp_analyzer import webapp_analyzer from .tool_wrappers import with_tool_events, wrap_tool_with_events +from .avfs import avfs_mount, avfs_umount, avfs_chdir, avfs_list, avfs_read, avfs_write, avfs_grep __all__ = [ @@ -41,7 +42,15 @@ "grep_session_logs", # web app analyzer "webapp_analyzer", + # AVFS + "avfs_mount", + "avfs_umount", + "avfs_chdir", + "avfs_list", + "avfs_read", + "avfs_write", + "avfs_grep", # Tool wrappers "with_tool_events", "wrap_tool_with_events", -] \ No newline at end of file +] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py new file mode 100644 index 0000000..0cea425 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py @@ -0,0 +1,17 @@ +from .avfs import AVFS, avfs +from .list import avfs_chdir, avfs_list, avfs_mount, avfs_umount +from .read import avfs_grep, avfs_read +from .write import avfs_write, write_text + +__all__ = [ + "AVFS", + "avfs", + "avfs_mount", + "avfs_umount", + "avfs_chdir", + "avfs_list", + "avfs_read", + "avfs_write", + "avfs_grep", + "write_text", +] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/avfs.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/avfs.py new file mode 100644 index 0000000..15719e6 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/avfs.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +import threading + + +@dataclass(frozen=True) +class _AVFSState: + """Mounted filesystem view for one session.""" + + root: Path + cwd: PurePosixPath + + +class AVFS: + """Session-scoped virtual filesystem rooted in a host workspace directory.""" + + def __init__(self) -> None: + self._lock = threading.RLock() + self._global_states: dict[str, _AVFSState] = {} + self._session_states: dict[str, dict[str, _AVFSState]] = {} + + def mount( + self, + workspace_root: str | Path, + *, + directory: str = ".", + session_id: str | None = None, + workspace: str = "workspace", + ) -> Path: + """Register a host workspace directory and initialize the virtual working directory.""" + root = Path(workspace_root).expanduser().resolve() + if not root.exists(): + raise FileNotFoundError(f"Workspace root does not exist: {root}") + if not root.is_dir(): + raise NotADirectoryError(f"Workspace root is not a directory: {root}") + + cwd = self._normalize_virtual_path(directory) + cwd_path = self._host_path_from_virtual(root, cwd) + if not cwd_path.exists(): + raise FileNotFoundError(f"Initial AVFS directory does not exist: {self.format_virtual_path(cwd)}") + if not cwd_path.is_dir(): + raise NotADirectoryError(f"Initial AVFS directory is not a directory: {self.format_virtual_path(cwd)}") + + state = _AVFSState(root=root, cwd=cwd) + with self._lock: + if session_id: + self._session_states.setdefault(session_id, {})[workspace] = state + else: + self._global_states[workspace] = state + return root + + def umount(self, session_id: str | None = None, *, workspace: str = "workspace") -> None: + """Unmount a session workspace or the global workspace.""" + with self._lock: + if session_id: + states = self._session_states.get(session_id) + if states is None: + return + states.pop(workspace, None) + if not states: + self._session_states.pop(session_id, None) + else: + self._global_states.pop(workspace, None) + + def current_mount(self, session_id: str | None = None, *, workspace: str = "workspace") -> Path | None: + """Backward-compatible alias for the active workspace root.""" + return self.current_workspace_root(session_id=session_id, workspace=workspace) + + def current_workspace_root(self, session_id: str | None = None, *, workspace: str = "workspace") -> Path | None: + """Return the active workspace root for a session.""" + state = self._current_state(session_id=session_id, workspace=workspace) + return None if state is None else state.root + + def current_directory(self, session_id: str | None = None, *, workspace: str = "workspace") -> str: + """Return the current virtual working directory.""" + state = self._require_state(session_id=session_id, workspace=workspace) + return self.format_virtual_path(state.cwd) + + def chdir(self, path: str, session_id: str | None = None, *, workspace: str = "workspace") -> str: + """Change the virtual working directory for a mounted session.""" + with self._lock: + state = self._require_state(session_id=session_id, workspace=workspace) + target = self._resolve_virtual_path(path, cwd=state.cwd) + target_path = self._host_path_from_virtual(state.root, target) + if not target_path.exists(): + raise FileNotFoundError(f"Directory does not exist: {self.format_virtual_path(target)}") + if not target_path.is_dir(): + raise NotADirectoryError(f"Path is not a directory: {self.format_virtual_path(target)}") + + next_state = _AVFSState(root=state.root, cwd=target) + if session_id: + self._session_states.setdefault(session_id, {})[workspace] = next_state + else: + self._global_states[workspace] = next_state + return self.format_virtual_path(target) + + def resolve_virtual_path( + self, + path: str = ".", + *, + session_id: str | None = None, + workspace: str = "workspace", + ) -> PurePosixPath: + """Resolve a user path inside the mounted virtual namespace.""" + state = self._require_state(session_id=session_id, workspace=workspace) + return self._resolve_virtual_path(path, cwd=state.cwd) + + def resolve(self, path: str = ".", session_id: str | None = None, *, workspace: str = "workspace") -> Path: + """Resolve a path inside the workspace root using virtual cwd semantics.""" + state = self._require_state(session_id=session_id, workspace=workspace) + virtual_path = self._resolve_virtual_path(path, cwd=state.cwd) + return self._host_path_from_virtual(state.root, virtual_path) + + def to_virtual_path( + self, + host_path: str | Path, + session_id: str | None = None, + *, + workspace: str = "workspace", + ) -> PurePosixPath: + """Convert a resolved host path back into a virtual path.""" + state = self._require_state(session_id=session_id, workspace=workspace) + resolved = Path(host_path).resolve(strict=False) + try: + relative = resolved.relative_to(state.root) + except ValueError as exc: + raise ValueError(f"Path is outside workspace root: {host_path}") from exc + return PurePosixPath(*relative.parts) if relative.parts else PurePosixPath(".") + + @staticmethod + def format_virtual_path(path: PurePosixPath) -> str: + """Render a virtual path with a stable root-relative representation.""" + return "/" if path == PurePosixPath(".") else f"/{path.as_posix()}" + + def _current_state(self, session_id: str | None = None, *, workspace: str = "workspace") -> _AVFSState | None: + with self._lock: + if session_id and session_id in self._session_states: + return self._session_states[session_id].get(workspace) + return self._global_states.get(workspace) + + def _require_state(self, session_id: str | None = None, *, workspace: str = "workspace") -> _AVFSState: + state = self._current_state(session_id=session_id, workspace=workspace) + if state is None: + raise RuntimeError(f"AVFS workspace '{workspace}' is not mounted. Call avfs_mount first.") + return state + + def _resolve_virtual_path(self, path: str, *, cwd: PurePosixPath) -> PurePosixPath: + requested = path or "." + requested_path = PurePosixPath(requested) + if requested_path.is_absolute(): + return self._normalize_virtual_path(requested) + return self._normalize_virtual_path(str(cwd / requested_path)) + + def _normalize_virtual_path(self, path: str) -> PurePosixPath: + pure_path = PurePosixPath(path or ".") + parts = pure_path.parts[1:] if pure_path.is_absolute() else pure_path.parts + + normalized_parts: list[str] = [] + for part in parts: + if part in ("", "."): + continue + if part == "..": + if not normalized_parts: + raise ValueError(f"Path escapes workspace root: {path}") + normalized_parts.pop() + continue + normalized_parts.append(part) + + return PurePosixPath(*normalized_parts) if normalized_parts else PurePosixPath(".") + + def _host_path_from_virtual(self, root: Path, virtual_path: PurePosixPath) -> Path: + candidate = root if virtual_path == PurePosixPath(".") else root.joinpath(*virtual_path.parts) + resolved = candidate.resolve(strict=False) + try: + resolved.relative_to(root) + except ValueError as exc: + raise ValueError(f"Path escapes workspace root: {self.format_virtual_path(virtual_path)}") from exc + return resolved + + +avfs = AVFS() diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py new file mode 100644 index 0000000..26131b2 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import os +from pathlib import Path +from pydantic_ai import RunContext + +from deadend_agent.tools.avfs.avfs import avfs +from deadend_agent.tools.tool_wrappers import with_tool_events + + +def _session_id_from_ctx(ctx: RunContext[object]) -> str | None: + deps = getattr(ctx, "deps", None) + if deps is None: + return None + value = getattr(deps, "session_id", None) + return str(value) if value is not None else None + + +@with_tool_events("avfs_mount") +async def avfs_mount( + ctx: RunContext[object], + workspace_root: str, + directory: str = ".", + workspace: str = "workspace", +) -> str: + """Register a workspace root and initialize the virtual working directory.""" + session_id = _session_id_from_ctx(ctx) + mounted = avfs.mount( + workspace_root=workspace_root, + directory=directory, + session_id=session_id, + workspace=workspace, + ) + return f"AVFS workspace '{workspace}': {mounted} (cwd={avfs.current_directory(session_id=session_id, workspace=workspace)})" + + +@with_tool_events("avfs_umount") +async def avfs_umount( + ctx: RunContext[object], + workspace: str = "workspace", +) -> str: + """Unmount AVFS for current session.""" + avfs.umount(session_id=_session_id_from_ctx(ctx), workspace=workspace) + return f"Unmounted AVFS workspace '{workspace}'." + + +@with_tool_events("avfs_chdir") +async def avfs_chdir( + ctx: RunContext[object], + path: str, + workspace: str = "workspace", +) -> str: + """Change the virtual working directory inside the mounted AVFS root.""" + directory = avfs.chdir(path, session_id=_session_id_from_ctx(ctx), workspace=workspace) + return f"Changed AVFS directory to {directory}" + + +@with_tool_events("avfs_list") +async def avfs_list( + ctx: RunContext[object], + path: str = ".", + recursive: bool = False, + include_hidden: bool = False, + max_entries: int = 200, + workspace: str = "workspace", +) -> list[dict[str, str | int | bool]]: + """List files and directories inside the current workspace root.""" + session_id = _session_id_from_ctx(ctx) + target = avfs.resolve(path, session_id=session_id, workspace=workspace) + if not target.exists(): + raise FileNotFoundError(f"Path does not exist: {path}") + if not target.is_dir(): + raise NotADirectoryError(f"Path is not a directory: {path}") + + workspace_root = avfs.current_workspace_root(session_id=session_id, workspace=workspace) + assert workspace_root is not None + + items: list[dict[str, str | int | bool]] = [] + if max_entries <= 0: + return items + + if recursive: + for current_root, dirnames, filenames in os.walk(target, topdown=True, followlinks=False): + dirnames.sort() + filenames.sort() + if not include_hidden: + dirnames[:] = [name for name in dirnames if not name.startswith(".")] + + current_root_path = Path(current_root) + for name in dirnames + filenames: + if not include_hidden and name.startswith("."): + continue + node = current_root_path / name + rel = node.relative_to(workspace_root).as_posix() + items.append( + { + "path": rel, + "type": "directory" if node.is_dir() else "file", + "size_bytes": node.stat().st_size if node.is_file() else 0, + "is_hidden": any(part.startswith(".") for part in Path(rel).parts), + } + ) + if len(items) >= max_entries: + return items + else: + for node in sorted(target.iterdir(), key=lambda entry: entry.name): + rel = node.relative_to(workspace_root).as_posix() + if not include_hidden and any(part.startswith(".") for part in Path(rel).parts): + continue + items.append( + { + "path": rel, + "type": "directory" if node.is_dir() else "file", + "size_bytes": node.stat().st_size if node.is_file() else 0, + "is_hidden": any(part.startswith(".") for part in Path(rel).parts), + } + ) + if len(items) >= max_entries: + return items + return items diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py new file mode 100644 index 0000000..a043681 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any +from pydantic_ai import RunContext + +from deadend_agent.tools.avfs.avfs import avfs +from deadend_agent.tools.tool_wrappers import with_tool_events + + +def _session_id_from_ctx(ctx: RunContext[object]) -> str | None: + deps = getattr(ctx, "deps", None) + if deps is None: + return None + value = getattr(deps, "session_id", None) + return str(value) if value is not None else None + + +@with_tool_events("avfs_read") +async def avfs_read( + ctx: RunContext[object], + path: str, + start_line: int = 1, + end_line: int | None = None, + max_chars: int = 100_000, + workspace: str = "workspace", +) -> str: + """Read a text file inside the current workspace root with optional 1-based line slicing.""" + session_id = _session_id_from_ctx(ctx) + target = avfs.resolve(path, session_id=session_id, workspace=workspace) + if not target.exists() or not target.is_file(): + raise FileNotFoundError(f"File not found: {path}") + if start_line < 1: + raise ValueError("start_line must be >= 1") + if end_line is not None and end_line < start_line: + raise ValueError("end_line must be >= start_line") + if max_chars <= 0: + raise ValueError("max_chars must be > 0") + + chunks: list[str] = [] + total_chars = 0 + truncated = False + try: + with open(target, "r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if line_number < start_line: + continue + if end_line is not None and line_number > end_line: + break + + remaining = max_chars - total_chars + if remaining <= 0: + truncated = True + break + + if len(line) <= remaining: + chunks.append(line) + total_chars += len(line) + continue + + chunks.append(line[:remaining]) + total_chars += remaining + truncated = True + break + except UnicodeDecodeError as exc: + raise ValueError(f"File is not valid UTF-8 text: {path}") from exc + + result = "".join(chunks) + if truncated: + suffix = "\n...[truncated]" + if len(result) + len(suffix) <= max_chars: + result = f"{result}{suffix}" + return result + + +@with_tool_events("avfs_grep") +async def avfs_grep( + ctx: RunContext[object], + pattern: str, + path: str = ".", + max_results: int = 50, + case_sensitive: bool = False, + include_hidden: bool = False, + workspace: str = "workspace", +) -> list[dict[str, str | int]]: + """Search files in the current workspace root using a regex pattern.""" + session_id = _session_id_from_ctx(ctx) + target = avfs.resolve(path, session_id=session_id, workspace=workspace) + if not target.exists(): + raise FileNotFoundError(f"Path not found: {path}") + if max_results <= 0: + return [] + + Ripgrepy = _load_ripgrepy() + search = Ripgrepy(pattern, str(target)).json().no_config().no_messages().no_ignore().m(max_results) + if not case_sensitive: + search = search.ignore_case() + if include_hidden: + search = search.hidden() + + raw_matches = search.run().as_dict + matches: list[dict[str, str | int]] = [] + for raw_match in raw_matches: + if raw_match.get("type") != "match": + continue + data = raw_match.get("data", {}) + host_path = _extract_text_field(data.get("path", {})) + if not host_path: + continue + line_text = _extract_text_field(data.get("lines", {})).rstrip("\n") + line_number = int(data.get("line_number", 0)) + submatches = data.get("submatches", []) + + if not submatches: + matches.append( + { + "path": avfs.to_virtual_path(host_path, session_id=session_id, workspace=workspace).as_posix(), + "line_number": line_number, + "match": "", + "context": line_text[:240], + } + ) + else: + for submatch in submatches: + matched_text = _extract_text_field(submatch.get("match", {})) + matches.append( + { + "path": avfs.to_virtual_path(host_path, session_id=session_id, workspace=workspace).as_posix(), + "line_number": line_number, + "match": matched_text, + "context": line_text[:240], + } + ) + if len(matches) >= max_results: + return matches + if len(matches) >= max_results: + return matches + return matches + + +def _load_ripgrepy() -> Any: + try: + from ripgrepy import Ripgrepy + except ModuleNotFoundError as exc: + raise RuntimeError("ripgrepy is required for avfs_grep but is not installed.") from exc + return Ripgrepy + + +def _extract_text_field(value: Any) -> str: + if isinstance(value, dict): + if "text" in value and value["text"] is not None: + return str(value["text"]) + if "bytes" in value and value["bytes"] is not None: + return str(value["bytes"]) + if value is None: + return "" + return str(value) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py new file mode 100644 index 0000000..fc4d9b7 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from pydantic_ai import RunContext + +from deadend_agent.tools.avfs.avfs import avfs +from deadend_agent.tools.tool_wrappers import with_tool_events + + +def _session_id_from_ctx(ctx: RunContext[object]) -> str | None: + deps = getattr(ctx, "deps", None) + if deps is None: + return None + value = getattr(deps, "session_id", None) + return str(value) if value is not None else None + + +def write_text( + path: str, + content: str, + *, + session_id: str | None, + workspace: str = "workspace", + append: bool = False, +) -> str: + """Write content to a virtual path without requiring a RunContext.""" + target = avfs.resolve(path, session_id=session_id, workspace=workspace) + if target.exists() and target.is_dir(): + raise IsADirectoryError(f"Cannot write to directory: {path}") + target.parent.mkdir(parents=True, exist_ok=True) + mode = "a" if append else "w" + with open(target, mode, encoding="utf-8") as handle: + handle.write(content) + virtual_path = avfs.resolve_virtual_path(path, session_id=session_id, workspace=workspace) + return f"Wrote {len(content.encode('utf-8'))} bytes to {avfs.format_virtual_path(virtual_path)}" + + +@with_tool_events("avfs_write") +async def avfs_write( + ctx: RunContext[object], + path: str, + content: str, + append: bool = False, + workspace: str = "workspace", +) -> str: + """Write content to a file under the current workspace root.""" + session_id = _session_id_from_ctx(ctx) + return write_text( + path, + content, + session_id=session_id, + workspace=workspace, + append=append, + ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/__init__.py index d0fc880..b24f720 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/__init__.py @@ -76,7 +76,7 @@ async def read_auth_storage(ctx: str) -> str: @with_tool_events("run_python_file") async def run_python_file( - ctx: RunContext[str], + ctx: RunContext[Any], code: str, filename: str, packages: list[str] @@ -109,9 +109,12 @@ async def run_python_file( file_path.write_text(code, encoding="utf-8") print(code) - # Get session_id from context deps (passed from agent), or generate one if not provided - # ctx.deps is the session_id string passed from PythonInterpreterAgent - session_id = ctx.deps if ctx.deps and isinstance(ctx.deps, str) else f"session_{id(file_path)}" + deps = getattr(ctx, "deps", None) + if isinstance(deps, str): + session_id = deps + else: + session_id = getattr(deps, "session_id", None) if deps is not None else None + session_id = session_id or f"session_{id(file_path)}" # Initializing the PythonInterpreter # Convert cache_dir Path to string for the directory parameter diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/python_interpreter.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/python_interpreter.py index 0decf58..6a63eaf 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/python_interpreter.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/python_interpreter/python_interpreter.py @@ -4,198 +4,65 @@ """Python interpreter tool for executing Python code in sandboxed environments. -This module provides functionality to execute Python code safely within -sandboxed environments, enabling AI agents to run Python scripts and -code snippets for security research and analysis tasks. - -The python sandbox is a WebAssembly server that is ran from a binary : `python-sandbox-tool` -This binary is compiled from : https://github.com/xoxruns/simple-python-interpreter-sandbox -and will be intergrated to the whole project in the future. +This adapter keeps the original ``PythonInterpreter`` interface used by the +agent code while routing execution to the stdio worker pool implementation from +``python_sandbox_client``. """ -import asyncio -from enum import Enum, unique -from asyncio.subprocess import PIPE, Process from pathlib import Path from typing import Any -import aiohttp -ENDPOINT_PYTHON_SANDBOX="http://127.0.0.1:45555" -PYTHON_SANDBOX_NAME="python-sandbox-tool-linux" +from python_sandbox_client import SandboxPool class PythonInterpreterNotFoundException(FileNotFoundError): - """Raised when the sandbox binary cannot be found locally.""" - -@unique -class CommandsInterpreter(str, Enum): - """HTTP endpoints exposed by the sandboxed Python interpreter service.""" - INSTALL_PACKAGES = f"{ENDPOINT_PYTHON_SANDBOX}/installpackages" - RUN_SCRIPT = f"{ENDPOINT_PYTHON_SANDBOX}/runscript" - CHECK_PACKAGES = f"{ENDPOINT_PYTHON_SANDBOX}/checkpackages" - SET_DIRECTORY = f"{ENDPOINT_PYTHON_SANDBOX}/setdirectory" + """Kept for backward compatibility with existing imports.""" class PythonInterpreter: - """Manage lifecycle of the sandboxed Python interpreter and issue HTTP commands. - - Responsibilities: - - Ensure the sandbox binary exists locally (download on first use). - - Start and stop the sandbox process. - - Send JSON requests (set directory, install packages, run scripts). - """ + """Manage lifecycle of the sandboxed Python interpreter client.""" session_id: str | None directory: str - pid: Process | None = None + pool: SandboxPool | None = None def __init__(self, session_id: str | None, directory: str) -> None: self.session_id = session_id - self.directory = directory - self.cache_python_dir = Path.home() / ".cache" / "deadend" / "python" - self.cache_python_dir.mkdir(parents=True, exist_ok=True) - self.cache_python_sandbox = self.cache_python_dir / PYTHON_SANDBOX_NAME + self.directory = str(Path(directory).resolve()) async def initialize(self): - """Ensure the sandbox binary exists, start the process, set working directory. - - Downloads the binary if missing, spawns the process if not already running, - then calls the sandbox to set the working directory. - - Returns: - Any: JSON response from the sandbox for the set-directory request. - """ - - # Downloads the python-sandbox-tool binary to cache if it doesn't exist - # This is a lot of context managers for a simple download. - # We need to add a checksum verification here - if not self.cache_python_sandbox.exists(): - raise PythonInterpreterNotFoundException( - f"Python sandbox not found at {self.cache_python_sandbox}. " - "Download it first via deadend_agent.core.download_python_sandbox()." - ) - - # and starts the process - if self.pid and self.pid.returncode is None: - return - - self.pid = await asyncio.create_subprocess_exec( - program=str(self.cache_python_sandbox), - stdout=PIPE, stderr=PIPE, - cwd=self.directory - ) - - # Setting the directory - # NOTE: the sandbox HTTP server may not be ready immediately after the - # process starts, so we add a small retry loop here to avoid transient - # "connection refused" errors that surface as tool failures like: - # "Error executing tool: CommandsInterpreter.SET_DIRECTORY". - last_exc: Exception | None = None - for attempt in range(5): - try: - resp = await self._send_instruction_post( - command=CommandsInterpreter.SET_DIRECTORY, - key="directory", - data=self.directory, - ) - print(resp) - return resp - except aiohttp.ClientError as exc: # type: ignore[attr-defined] - last_exc = exc - # Back off slightly between attempts to give the server time to boot - await asyncio.sleep(0.2 * (attempt + 1)) - - # If we got here, all attempts failed – raise a clear, high-level error - raise RuntimeError( - f"Failed to reach Python sandbox at {str(CommandsInterpreter.SET_DIRECTORY)} " - f"after 5 attempts. Last error: {last_exc!r}" - ) + """Start sandbox pool and bind to the configured working directory.""" + if self.pool is not None: + return {"current_directory": self.directory} + self.pool = SandboxPool(directory=self.directory, workers=1) + await self.pool.__aenter__() + return {"current_directory": self.directory} async def load_packages(self, packages: list[str]): - """Request package installation inside the sandbox. - - Args: - packages: List of package specifiers (e.g., ["requests==2.32.3", "numpy"]). - - Returns: - Any: JSON response from the sandbox. - """ - # Loads the packages needed for the file - if self.pid is None: + """Request package installation inside the sandbox.""" + if self.pool is None: raise RuntimeError("Interpreter not initialized. Call initialize() first.") - - return await self._send_instruction_post( - command=CommandsInterpreter.INSTALL_PACKAGES, - key="packages", - data=packages - ) + if not packages: + return {"results": {}} + return {"results": await self.pool.install_packages(packages)} async def run_file(self, filename: str, _session_id: str | None = None): - """Execute a Python file within the configured working directory. - - Args: - filename: Relative path to the file to execute. - session_id: Optional override of the current session identifier. - - Returns: - Any: JSON response from the sandbox with execution result. - """ - # Run a file present in the directory specified - return await self._send_instruction_post( - command=CommandsInterpreter.RUN_SCRIPT, - key="filename", - data=filename - ) + """Execute a Python file within the configured working directory.""" + if self.pool is None: + raise RuntimeError("Interpreter not initialized. Call initialize() first.") + execution = await self.pool.run_script(filename) + return { + "result": execution.result, + "stdout": execution.stdout, + "stderr": execution.stderr, + } async def run_code(self, code: str): """Execute inline Python code in the sandbox (not implemented).""" raise NotImplementedError - async def _send_instruction_post(self, command: str | CommandsInterpreter, key: str, data: Any): - """Send a JSON POST request to the sandbox. - - Args: - command: Target URL (an entry from `CommandsInterpreter`). - key: JSON key for the payload. - data: JSON value to send under `key`. - - Returns: - Any: Parsed JSON response. - """ - # Explicitly convert Enum to string for aiohttp compatibility - # Using .value for Enums ensures we get the actual string value - # This is necessary for Python 3.10 compatibility (StrEnum is 3.11+) - if isinstance(command, Enum): - url = command.value - else: - url = str(command) - async with aiohttp.ClientSession() as session: - async with session.post(url=url, json={key: data}) as resp: - return await resp.json() - async def shutdown(self, timeout: float = 5.0): - """Gracefully terminate the sandbox process. - - Sends SIGTERM and waits up to `timeout`. If the process does not exit, - sends SIGKILL and waits for termination. - - Args: - timeout: Seconds to wait after terminate before force-killing. - """ - if not self.pid: + """Gracefully terminate worker pool.""" + del timeout # maintained for API compatibility + if self.pool is None: return - if self.pid.returncode is not None: - self.pid = None - return - try: - self.pid.terminate() - except ProcessLookupError: - pass - try: - await asyncio.wait_for(self.pid.wait(), timeout=timeout) - except asyncio.TimeoutError: - try: - self.pid.kill() - except ProcessLookupError: - pass - await self.pid.wait() - finally: - self.pid = None + await self.pool.aclose() + self.pool = None diff --git a/deadend_cli/deadend_agent/src/deadend_agent/utils/structures.py b/deadend_cli/deadend_agent/src/deadend_agent/utils/structures.py index 51688d2..b46153e 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/utils/structures.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/utils/structures.py @@ -114,8 +114,14 @@ class ShellDeps: Attributes: shell_runner: ShellRunner instance for command execution + session_id: Session identifier used by session-scoped tools such as AVFS + workspace_root: Optional host workspace mounted into AVFS for this session """ shell_runner: ShellRunner + session_id: uuid.UUID + workspace_root: str | None = None + memory_workspace_root: str | None = None + memory_context: str = "" @dataclass @@ -139,6 +145,8 @@ class RequesterDeps: target: str session_id: uuid.UUID embedding_session_id: uuid.UUID | None = None + memory_workspace_root: str | None = None + memory_context: str = "" @dataclass class WebappreconDeps: @@ -162,6 +170,16 @@ class WebappreconDeps: shell_runner: ShellRunner session_id: uuid.UUID embedding_session_id: uuid.UUID | None = None + memory_workspace_root: str | None = None + memory_context: str = "" + +@dataclass +class MemoryWorkspaceDeps: + """Dependencies shared by agents that need access to the memory workspace.""" + + session_id: str + memory_workspace_root: str | None = None + memory_context: str = "" @dataclass class RagDeps: diff --git a/deadend_cli/deadend_agent/tests/rlm/test_avfs.py b/deadend_cli/deadend_agent/tests/rlm/test_avfs.py new file mode 100644 index 0000000..14712ea --- /dev/null +++ b/deadend_cli/deadend_agent/tests/rlm/test_avfs.py @@ -0,0 +1,254 @@ +import asyncio +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "src" / "deadend_agent" + +if "deadend_agent" not in sys.modules: + package = types.ModuleType("deadend_agent") + package.__path__ = [str(PACKAGE_ROOT)] + sys.modules["deadend_agent"] = package + +if "deadend_agent.tools" not in sys.modules: + tools_package = types.ModuleType("deadend_agent.tools") + tools_package.__path__ = [str(PACKAGE_ROOT / "tools")] + sys.modules["deadend_agent.tools"] = tools_package + +if "deadend_agent.tools.avfs" not in sys.modules: + avfs_package = types.ModuleType("deadend_agent.tools.avfs") + avfs_package.__path__ = [str(PACKAGE_ROOT / "tools" / "avfs")] + sys.modules["deadend_agent.tools.avfs"] = avfs_package + +if "pydantic_ai" not in sys.modules: + pydantic_ai_module = types.ModuleType("pydantic_ai") + + class RunContext: + def __init__(self, deps=None): + self.deps = deps + + def __class_getitem__(cls, _item): + return cls + + pydantic_ai_module.RunContext = RunContext + sys.modules["pydantic_ai"] = pydantic_ai_module + +from deadend_agent.tools.avfs.avfs import AVFS, avfs +from deadend_agent.tools.avfs.list import avfs_mount, avfs_umount +from deadend_agent.tools.avfs.read import avfs_grep +from deadend_agent.tools.avfs.write import avfs_write, write_text + + +def test_avfs_mount_and_resolve(tmp_path): + root = tmp_path / "workspace" + root.mkdir() + (root / "notes.txt").write_text("hello", encoding="utf-8") + + fs = AVFS() + fs.mount(root) + + resolved = fs.resolve("notes.txt") + assert resolved == (root / "notes.txt").resolve() + assert fs.current_directory() == "/" + + +def test_avfs_blocks_path_escape(tmp_path): + root = tmp_path / "workspace" + root.mkdir() + fs = AVFS() + fs.mount(root) + + with pytest.raises(ValueError): + fs.resolve("../secret.txt") + + +def test_avfs_resolves_relative_to_virtual_directory(tmp_path): + root = tmp_path / "workspace" + nested = root / "src" / "pkg" + nested.mkdir(parents=True) + target = nested / "module.py" + target.write_text("print('ok')\n", encoding="utf-8") + + fs = AVFS() + fs.mount(root, directory="src") + + assert fs.current_directory() == "/src" + assert fs.resolve("pkg/module.py") == target.resolve() + assert fs.resolve("./pkg/../pkg/module.py") == target.resolve() + + +def test_avfs_mount_rejects_symlink_escape(tmp_path): + root = tmp_path / "workspace" + outside = tmp_path / "outside" + root.mkdir() + outside.mkdir() + (root / "jump").symlink_to(outside, target_is_directory=True) + + fs = AVFS() + fs.mount(root) + + with pytest.raises(ValueError): + fs.resolve("jump/secret.txt") + + +def test_avfs_can_change_virtual_directory(tmp_path): + root = tmp_path / "workspace" + nested = root / "docs" + nested.mkdir(parents=True) + + fs = AVFS() + fs.mount(root) + + changed = fs.chdir("docs") + + assert changed == "/docs" + assert fs.current_directory() == "/docs" + + +def test_avfs_tracks_session_specific_state(tmp_path): + root = tmp_path / "workspace" + (root / "alpha").mkdir(parents=True) + (root / "beta").mkdir(parents=True) + + fs = AVFS() + fs.mount(root, directory="alpha", session_id="session-a") + fs.mount(root, directory="beta", session_id="session-b") + + assert fs.current_directory(session_id="session-a") == "/alpha" + assert fs.current_directory(session_id="session-b") == "/beta" + + +def test_avfs_write_updates_host_file(tmp_path): + workspace_root = tmp_path / "workspace" + workspace_root.mkdir() + ctx = SimpleNamespace(deps=SimpleNamespace(session_id="session-write")) + + async def run_test() -> None: + await avfs_mount(ctx, workspace_root=str(workspace_root)) + try: + result = await avfs_write(ctx, "notes.txt", "alpha") + assert "Wrote" in result + assert (workspace_root / "notes.txt").read_text(encoding="utf-8") == "alpha" + + await avfs_write(ctx, "notes.txt", "\nbeta", append=True) + assert (workspace_root / "notes.txt").read_text(encoding="utf-8") == "alpha\nbeta" + finally: + await avfs_umount(ctx) + + asyncio.run(run_test()) + + +def test_write_text_updates_named_memory_workspace(tmp_path): + memory_root = tmp_path / "agents" / "local-agent" / "memory" + memory_root.mkdir(parents=True) + + avfs.mount(memory_root, session_id="session-memory", workspace="memory") + try: + result = write_text( + "summaries/requester.md", + "alpha", + session_id="session-memory", + workspace="memory", + ) + + assert "Wrote" in result + assert (memory_root / "summaries" / "requester.md").read_text(encoding="utf-8") == "alpha" + finally: + avfs.umount(session_id="session-memory", workspace="memory") + + +def test_avfs_grep_uses_ripgrepy_and_returns_virtual_paths(tmp_path): + workspace_root = tmp_path / "workspace" + workspace_root.mkdir() + target = workspace_root / "notes.txt" + target.write_text("alpha\n", encoding="utf-8") + ctx = SimpleNamespace(deps=SimpleNamespace(session_id="session-grep")) + + class FakeRunResult: + @property + def as_dict(self): + return [ + { + "type": "match", + "data": { + "path": {"text": str(target)}, + "line_number": 1, + "lines": {"text": "alpha\n"}, + "submatches": [{"match": {"text": "alpha"}}], + }, + } + ] + + class FakeRipgrepy: + calls: list[tuple[str, object]] = [] + + def __init__(self, pattern: str, path: str) -> None: + self.calls.append(("init", (pattern, path))) + + def json(self): + self.calls.append(("json", None)) + return self + + def no_config(self): + self.calls.append(("no_config", None)) + return self + + def no_messages(self): + self.calls.append(("no_messages", None)) + return self + + def no_ignore(self): + self.calls.append(("no_ignore", None)) + return self + + def ignore_case(self): + self.calls.append(("ignore_case", None)) + return self + + def hidden(self): + self.calls.append(("hidden", None)) + return self + + def m(self, max_results): + self.calls.append(("m", max_results)) + return self + + def run(self): + self.calls.append(("run", None)) + return FakeRunResult() + + async def run_test() -> None: + await avfs_mount(ctx, workspace_root=str(workspace_root)) + try: + import sys + + sys.modules["ripgrepy"] = SimpleNamespace(Ripgrepy=FakeRipgrepy) + matches = await avfs_grep(ctx, "alpha", include_hidden=True) + + assert matches == [ + { + "path": "notes.txt", + "line_number": 1, + "match": "alpha", + "context": "alpha", + } + ] + assert FakeRipgrepy.calls == [ + ("init", ("alpha", str(workspace_root))), + ("json", None), + ("no_config", None), + ("no_messages", None), + ("no_ignore", None), + ("m", 50), + ("ignore_case", None), + ("hidden", None), + ("run", None), + ] + finally: + sys.modules.pop("ripgrepy", None) + await avfs_umount(ctx) + + asyncio.run(run_test()) diff --git a/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py b/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py new file mode 100644 index 0000000..0dc32cb --- /dev/null +++ b/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py @@ -0,0 +1,450 @@ +import asyncio +import sys +import types +from dataclasses import dataclass +from pathlib import Path +from uuid import UUID, uuid4 + + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "src" / "deadend_agent" + +if "deadend_agent" not in sys.modules: + package = types.ModuleType("deadend_agent") + package.__path__ = [str(PACKAGE_ROOT)] + sys.modules["deadend_agent"] = package + +if "deadend_agent.tools" not in sys.modules: + tools_package = types.ModuleType("deadend_agent.tools") + tools_package.__path__ = [str(PACKAGE_ROOT / "tools")] + sys.modules["deadend_agent.tools"] = tools_package + +for module_name in [ + "deadend_agent.tools.avfs", + "deadend_agent.tools.avfs.avfs", + "deadend_agent.tools.avfs.list", + "deadend_agent.tools.avfs.read", + "deadend_agent.tools.avfs.write", +]: + sys.modules.pop(module_name, None) + +if "deadend_agent.agents" not in sys.modules: + agents_package = types.ModuleType("deadend_agent.agents") + agents_package.__path__ = [str(PACKAGE_ROOT / "agents")] + sys.modules["deadend_agent.agents"] = agents_package + +if "deadend_agent.agents.generic_agents" not in sys.modules: + generic_agents_package = types.ModuleType("deadend_agent.agents.generic_agents") + generic_agents_package.__path__ = [str(PACKAGE_ROOT / "agents" / "generic_agents")] + sys.modules["deadend_agent.agents.generic_agents"] = generic_agents_package + +pydantic_ai_module = sys.modules.get("pydantic_ai") +if pydantic_ai_module is None: + pydantic_ai_module = types.ModuleType("pydantic_ai") + sys.modules["pydantic_ai"] = pydantic_ai_module + + +class Tool: + def __init__(self, function, **_kwargs): + self.function = function + + +class DeferredToolRequests: + pass + + +class DeferredToolResults: + pass + + +class RunContext: + def __init__(self, deps=None): + self.deps = deps + + def __class_getitem__(cls, _item): + return cls + + +class RunUsage: + pass + + +class UsageLimits: + def __init__(self, *args, **kwargs): + pass + + +pydantic_ai_module.Tool = getattr(pydantic_ai_module, "Tool", Tool) +pydantic_ai_module.DeferredToolRequests = getattr(pydantic_ai_module, "DeferredToolRequests", DeferredToolRequests) +pydantic_ai_module.DeferredToolResults = getattr(pydantic_ai_module, "DeferredToolResults", DeferredToolResults) +pydantic_ai_module.RunContext = getattr(pydantic_ai_module, "RunContext", RunContext) +pydantic_ai_module.RunUsage = getattr(pydantic_ai_module, "RunUsage", RunUsage) +pydantic_ai_module.UsageLimits = getattr(pydantic_ai_module, "UsageLimits", UsageLimits) + +if "pydantic_ai.usage" not in sys.modules: + usage_module = types.ModuleType("pydantic_ai.usage") + + class RunUsage: + pass + + class UsageLimits: + def __init__(self, *args, **kwargs): + pass + + usage_module.RunUsage = RunUsage + usage_module.UsageLimits = UsageLimits + sys.modules["pydantic_ai.usage"] = usage_module + +if "pydantic" not in sys.modules: + pydantic_module = types.ModuleType("pydantic") + + class BaseModel: + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + def Field(*, default=None, default_factory=None, description=None): + if default_factory is not None: + return default_factory() + return default + + pydantic_module.BaseModel = BaseModel + pydantic_module.Field = Field + sys.modules["pydantic"] = pydantic_module + +if "deadend_agent.logging" not in sys.modules: + logging_module = types.ModuleType("deadend_agent.logging") + logging_module.logger = object() + sys.modules["deadend_agent.logging"] = logging_module + +if "deadend_agent.config.settings" not in sys.modules: + settings_module = types.ModuleType("deadend_agent.config.settings") + + class ModelSpec: + pass + + class Config: + agents_storage_root = str(Path("/tmp") / "deadend-agent-tests") + + @classmethod + def get_local_agent_id(cls): + return uuid4() + + settings_module.Config = Config + settings_module.ModelSpec = ModelSpec + sys.modules["deadend_agent.config.settings"] = settings_module + +if "deadend_agent.models.registry" not in sys.modules: + registry_module = types.ModuleType("deadend_agent.models.registry") + + class EmbedderClient: + pass + + registry_module.EmbedderClient = EmbedderClient + sys.modules["deadend_agent.models.registry"] = registry_module + +if "deadend_agent.embedders.code_indexer" not in sys.modules: + code_indexer_module = types.ModuleType("deadend_agent.embedders.code_indexer") + + class SourceCodeIndexer: + def __init__(self, *args, **kwargs): + pass + + code_indexer_module.SourceCodeIndexer = SourceCodeIndexer + sys.modules["deadend_agent.embedders.code_indexer"] = code_indexer_module + +if "deadend_agent.context" not in sys.modules: + context_module = types.ModuleType("deadend_agent.context") + + class ContextEngine: + def __init__(self, *args, **kwargs): + self.target = None + + def set_target(self, target): + self.target = target + + context_module.ContextEngine = ContextEngine + sys.modules["deadend_agent.context"] = context_module + +if "deadend_agent.rag.sqlite_connector" not in sys.modules: + rag_module = types.ModuleType("deadend_agent.rag.sqlite_connector") + + class SqliteRagConnector: + pass + + rag_module.SqliteRagConnector = SqliteRagConnector + sys.modules["deadend_agent.rag.sqlite_connector"] = rag_module + +if "deadend_agent.sandbox.sandbox" not in sys.modules: + sandbox_module = types.ModuleType("deadend_agent.sandbox.sandbox") + + class Sandbox: + pass + + sandbox_module.Sandbox = Sandbox + sys.modules["deadend_agent.sandbox.sandbox"] = sandbox_module + +for module_name, class_name in [ + ("deadend_agent.agents.reporter", "ReporterAgent"), + ("deadend_agent.agents.architecture", "ADaPTAgent"), + ("deadend_agent.agents.recon_threatmodel_agent", "ReconThreatModelAgent"), + ("deadend_agent.agents.exploit_web_agent", "PlannerExploitAgent"), +]: + if module_name not in sys.modules: + module = types.ModuleType(module_name) + module.__dict__[class_name] = type(class_name, (), {}) + sys.modules[module_name] = module + +if "deadend_agent.agents.components.executor" not in sys.modules: + executor_module = types.ModuleType("deadend_agent.agents.components.executor") + + class ResultEvent: + pass + + class AgentExecutor: + def __init__(self, *args, **kwargs): + self.dependencies = {} + self.memory_context = "" + self.auth_session_key = "" + + def set_dependencies(self, **kwargs): + self.dependencies = kwargs + + def set_memory_context(self, memory_context: str): + self.memory_context = memory_context + + def set_auth_session_key(self, auth_session_key: str): + self.auth_session_key = auth_session_key + + executor_module.ResultEvent = ResultEvent + executor_module.AgentExecutor = AgentExecutor + sys.modules["deadend_agent.agents.components.executor"] = executor_module + +if "deadend_agent.agents.components.planner" not in sys.modules: + planner_module = types.ModuleType("deadend_agent.agents.components.planner") + planner_module.Planner = type("Planner", (), {}) + planner_module.TaskNode = type("TaskNode", (), {}) + sys.modules["deadend_agent.agents.components.planner"] = planner_module + +if "deadend_agent.agents.components.validator" not in sys.modules: + validator_module = types.ModuleType("deadend_agent.agents.components.validator") + + class Validator: + def __init__(self, *args, **kwargs): + pass + + validator_module.Validator = Validator + sys.modules["deadend_agent.agents.components.validator"] = validator_module + +if "deadend_agent.utils.structures" not in sys.modules: + structures_module = types.ModuleType("deadend_agent.utils.structures") + + class ShellRunner: + def __init__(self, session: str, sandbox): + self.session = session + self.sandbox = sandbox + + @dataclass + class ShellDeps: + shell_runner: ShellRunner + session_id: str + workspace_root: str | None = None + memory_workspace_root: str | None = None + memory_context: str = "" + + @dataclass + class RequesterDeps: + embedder_client: object + rag: object + target: str + session_id: object + embedding_session_id: object | None = None + memory_workspace_root: str | None = None + memory_context: str = "" + + @dataclass + class WebappreconDeps: + embedder_client: object + rag: object + target: str + shell_runner: ShellRunner + session_id: object + embedding_session_id: object | None = None + memory_workspace_root: str | None = None + memory_context: str = "" + + @dataclass + class MemoryWorkspaceDeps: + session_id: str + memory_workspace_root: str | None = None + memory_context: str = "" + + structures_module.MemoryWorkspaceDeps = MemoryWorkspaceDeps + structures_module.RequesterDeps = RequesterDeps + structures_module.ShellDeps = ShellDeps + structures_module.ShellRunner = ShellRunner + structures_module.WebappreconDeps = WebappreconDeps + sys.modules["deadend_agent.utils.structures"] = structures_module + +if "deadend_agent.tools.browser_automation.http_parser" not in sys.modules: + http_parser_module = types.ModuleType("deadend_agent.tools.browser_automation.http_parser") + + def extract_host_port(target_host: str): + return "example.com", 443 + + http_parser_module.extract_host_port = extract_host_port + sys.modules["deadend_agent.tools.browser_automation.http_parser"] = http_parser_module + +if "deadend_prompts" not in sys.modules: + prompts_module = types.ModuleType("deadend_prompts") + + def render_agent_instructions(*args, **kwargs): + return "instructions" + + def render_tool_description(tool_name: str, **kwargs): + return tool_name + + prompts_module.render_agent_instructions = render_agent_instructions + prompts_module.render_tool_description = render_tool_description + sys.modules["deadend_prompts"] = prompts_module + +if "deadend_agent.agents.factory" not in sys.modules: + factory_module = types.ModuleType("deadend_agent.agents.factory") + + class AgentOutput: + pass + + class AgentRunner: + def __init__(self, name, model, instructions, deps_type, output_type, tools): + self.name = name + self.model = model + self.instructions = instructions + self.deps_type = deps_type + self.output_type = output_type + self.tools = tools + self.agent = types.SimpleNamespace(tools=tools) + + factory_module.AgentOutput = AgentOutput + factory_module.AgentRunner = AgentRunner + sys.modules["deadend_agent.agents.factory"] = factory_module + +if "deadend_agent.tools" in sys.modules: + tools_module = sys.modules["deadend_agent.tools"] + + async def sandboxed_shell_tool(*args, **kwargs): + return "ok" + + async def avfs_list(*args, **kwargs): + return [] + + async def avfs_read(*args, **kwargs): + return "" + + async def avfs_write(*args, **kwargs): + return "" + + async def avfs_grep(*args, **kwargs): + return [] + + tools_module.sandboxed_shell_tool = sandboxed_shell_tool + tools_module.avfs_list = avfs_list + tools_module.avfs_read = avfs_read + tools_module.avfs_write = avfs_write + tools_module.avfs_grep = avfs_grep + +from deadend_agent.deadend_agent import DeadEndAgent +from deadend_agent.tools.avfs import avfs +from deadend_agent.agents.generic_agents.shell_agent import ShellAgent + + +def test_deadend_agent_mounts_workspace_root_for_session(tmp_path): + workspace_root = tmp_path / "workspace" + workspace_root.mkdir() + storage_root = tmp_path / "agents" + embedding_session_id = uuid4() + + agent = DeadEndAgent( + session_id=uuid4(), + embedding_session_id=embedding_session_id, + model=object(), + available_agents={}, + workspace_root=str(workspace_root), + agents_storage_root=str(storage_root), + local_agent_id=UUID("11111111-1111-1111-1111-111111111111"), + ) + agent.target = "https://example.com" + agent.prepare_dependencies( + embedder_client=object(), + rag_connector=object(), + sandbox=object(), + target="https://example.com", + ) + + assert avfs.current_workspace_root(session_id=str(agent.agent_id)) == workspace_root.resolve() + assert avfs.current_workspace_root(session_id=str(agent.agent_id), workspace="memory") == ( + storage_root / "11111111-1111-1111-1111-111111111111" / str(embedding_session_id) / "memory" + ).resolve() + assert agent.shell_deps is not None + assert str(agent.agent_id) == "11111111-1111-1111-1111-111111111111" + assert str(agent.shell_deps.session_id) == str(agent.agent_id) + assert agent.shell_deps.workspace_root == str(workspace_root.resolve()) + assert agent.memory_workspace_root == str((storage_root / "11111111-1111-1111-1111-111111111111" / str(embedding_session_id) / "memory").resolve()) + assert agent.shell_deps.memory_workspace_root == str((storage_root / "11111111-1111-1111-1111-111111111111" / str(embedding_session_id) / "memory").resolve()) + + +def test_shell_agent_exposes_only_shell_tool(): + agent = ShellAgent( + model=object(), + deps_type=None, + target_information="target", + requires_approval=False, + ) + + tool_names = [tool.function.__name__ for tool in agent.tools] + assert tool_names == [ + "sandboxed_shell_tool", + ] + + +def test_deadend_agent_loads_memory_context_into_dependencies(tmp_path, monkeypatch): + workspace_root = tmp_path / "workspace" + workspace_root.mkdir() + storage_root = tmp_path / "agents" + embedding_session_id = uuid4() + + agent = DeadEndAgent( + session_id=uuid4(), + embedding_session_id=embedding_session_id, + model=object(), + available_agents={}, + workspace_root=str(workspace_root), + agents_storage_root=str(storage_root), + local_agent_id=UUID("11111111-1111-1111-1111-111111111111"), + ) + agent.target = "https://example.com" + agent.prepare_dependencies( + embedder_client=object(), + rag_connector=object(), + sandbox=object(), + target="https://example.com", + ) + + memory_module = sys.modules["deadend_agent.deadend_agent"] + + class FakeMemoryAgent: + def __init__(self, *args, **kwargs): + pass + + async def run(self, *args, **kwargs): + return types.SimpleNamespace(output="previous exploit worked") + + monkeypatch.setattr(memory_module, "MemoryAgent", FakeMemoryAgent) + + asyncio.run(agent._populate_memory_context("test objective")) + + assert agent.memory_context == "previous exploit worked" + assert agent.shell_deps is not None + assert agent.requester_deps is not None + assert agent.webapprecon_deps is not None + assert agent.shell_deps.memory_context == "previous exploit worked" + assert agent.requester_deps.memory_context == "previous exploit worked" + assert agent.webapprecon_deps.memory_context == "previous exploit worked" diff --git a/deadend_cli/deadend_agent/tests/rlm/test_memory.py b/deadend_cli/deadend_agent/tests/rlm/test_memory.py new file mode 100644 index 0000000..bd8a944 --- /dev/null +++ b/deadend_cli/deadend_agent/tests/rlm/test_memory.py @@ -0,0 +1,100 @@ +import json +import sys +import types +from pathlib import Path + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "src" / "deadend_agent" + +if "deadend_agent" not in sys.modules: + package = types.ModuleType("deadend_agent") + package.__path__ = [str(PACKAGE_ROOT)] + sys.modules["deadend_agent"] = package + +if "deadend_agent.context" not in sys.modules: + context_package = types.ModuleType("deadend_agent.context") + context_package.__path__ = [str(PACKAGE_ROOT / "context")] + sys.modules["deadend_agent.context"] = context_package + +if "deadend_agent.rlm" not in sys.modules: + rlm_package = types.ModuleType("deadend_agent.rlm") + rlm_package.__path__ = [str(PACKAGE_ROOT / "rlm")] + sys.modules["deadend_agent.rlm"] = rlm_package + +from deadend_agent.context.memory import MemoryHandler +from deadend_agent.rlm.compat import assess_python_sandbox_compatibility +from deadend_agent.rlm.memory import RLMFileMemory + + +def test_rlm_file_memory_indexes_markdown_and_json(tmp_path): + docs_dir = tmp_path / "session" + docs_dir.mkdir() + + markdown_path = docs_dir / "notes.md" + markdown_path.write_text( + "# Overview\n" + "System overview.\n\n" + "## Findings\n" + "- SQLi on /login\n" + "- Stored XSS on /profile\n", + encoding="utf-8", + ) + + json_path = docs_dir / "state.json" + json_path.write_text( + json.dumps( + { + "target": {"host": "example.com", "port": 443}, + "findings": [ + {"type": "sqli", "endpoint": "/login"}, + {"type": "xss", "endpoint": "/profile"}, + ], + } + ), + encoding="utf-8", + ) + + memory = RLMFileMemory(root=docs_dir) + files = memory.list_files() + + assert [item.path for item in files] == ["notes.md", "state.json"] + assert memory.read_md_section("notes.md", "Findings").startswith("## Findings") + assert memory.search_md_headings("find") == [ + { + "path": "notes.md", + "section_id": "findings", + "heading": "Findings", + "level": 2, + } + ] + + assert memory.json_get("state.json", "target.host") == "example.com" + assert memory.json_sample_array("state.json", "findings", 0, 1) == [ + {"type": "sqli", "endpoint": "/login"} + ] + + +def test_memory_handler_uses_session_root_and_describes_memory(tmp_path): + session_dir = tmp_path / "memory" / "sessions" / "test-session" + session_dir.mkdir(parents=True) + (session_dir / "requester.jsonl").write_text('{"event": "request"}\n', encoding="utf-8") + (session_dir / "summary.md").write_text("# Session\nUseful notes\n", encoding="utf-8") + + handler = MemoryHandler.for_session(session_key="test-session", base_dir=tmp_path / "memory" / "sessions") + + described = handler.describe_memory() + files = handler.list_files() + + assert "summary.md [markdown]" in described + assert "requester.jsonl [jsonl]" in described + assert [item.path for item in files] == ["requester.jsonl", "summary.md"] + + +def test_sandbox_compatibility_report_marks_current_backend_incompatible(): + report = assess_python_sandbox_compatibility() + + assert report.backend_name == "python-sandbox-tool" + assert report.compatible_for_full_rlm_repl is False + assert report.persistent_state is False + assert report.inline_code_execution is False + assert report.host_callback_support is False + assert len(report.blockers) >= 3 diff --git a/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py b/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py new file mode 100644 index 0000000..7b803b9 --- /dev/null +++ b/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py @@ -0,0 +1,72 @@ +import asyncio +import sys +import types +from pathlib import Path +from types import SimpleNamespace + + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "src" / "deadend_agent" + +if "deadend_agent" not in sys.modules: + package = types.ModuleType("deadend_agent") + package.__path__ = [str(PACKAGE_ROOT)] + sys.modules["deadend_agent"] = package + +if "deadend_agent.tools" not in sys.modules: + tools_package = types.ModuleType("deadend_agent.tools") + tools_package.__path__ = [str(PACKAGE_ROOT / "tools")] + sys.modules["deadend_agent.tools"] = tools_package + +if "pydantic_ai" not in sys.modules: + pydantic_ai_module = types.ModuleType("pydantic_ai") + + class RunContext: + def __init__(self, deps=None): + self.deps = deps + + def __class_getitem__(cls, _item): + return cls + + pydantic_ai_module.RunContext = RunContext + sys.modules["pydantic_ai"] = pydantic_ai_module + +from deadend_agent.tools.avfs.list import avfs_mount, avfs_umount +from deadend_agent.tools.avfs.read import avfs_read +from deadend_agent.tools.avfs.write import avfs_write + + +def test_avfs_write_updates_named_memory_workspace(tmp_path): + memory_root = tmp_path / "agents" / "local-agent" / "memory" + memory_root.mkdir(parents=True) + ctx = SimpleNamespace( + deps=SimpleNamespace( + session_id="memory-session", + memory_workspace_root=str(memory_root), + ) + ) + + async def run_test() -> None: + await avfs_mount(ctx, workspace_root=str(memory_root), workspace="memory") + try: + result = await avfs_write( + ctx, + "summaries/requester.md", + "alpha", + append=False, + workspace="memory", + ) + assert "Wrote" in result + assert (memory_root / "summaries" / "requester.md").read_text(encoding="utf-8") == "alpha" + + await avfs_write( + ctx, + "summaries/requester.md", + "\nbeta", + append=True, + workspace="memory", + ) + assert await avfs_read(ctx, "summaries/requester.md", workspace="memory") == "alpha\nbeta" + finally: + await avfs_umount(ctx, workspace="memory") + + asyncio.run(run_test()) diff --git a/deadend_cli/deadend_agent/tests/rlm/test_runner.py b/deadend_cli/deadend_agent/tests/rlm/test_runner.py new file mode 100644 index 0000000..a054f3f --- /dev/null +++ b/deadend_cli/deadend_agent/tests/rlm/test_runner.py @@ -0,0 +1,190 @@ +import json +import sys +import types +from pathlib import Path + + +PACKAGE_ROOT = Path(__file__).resolve().parents[2] / "src" / "deadend_agent" + +if "deadend_agent" not in sys.modules: + package = types.ModuleType("deadend_agent") + package.__path__ = [str(PACKAGE_ROOT)] + sys.modules["deadend_agent"] = package + +if "tenacity" not in sys.modules: + tenacity_module = types.ModuleType("tenacity") + + def retry(*args, **kwargs): + def decorator(func): + return func + return decorator + + def retry_if_exception_type(*args, **kwargs): + return None + + def stop_after_attempt(*args, **kwargs): + return None + + def wait_exponential(*args, **kwargs): + return None + + tenacity_module.retry = retry + tenacity_module.retry_if_exception_type = retry_if_exception_type + tenacity_module.stop_after_attempt = stop_after_attempt + tenacity_module.wait_exponential = wait_exponential + sys.modules["tenacity"] = tenacity_module + +if "deadend_agent.rlm" not in sys.modules: + rlm_package = types.ModuleType("deadend_agent.rlm") + rlm_package.__path__ = [str(PACKAGE_ROOT / "rlm")] + sys.modules["deadend_agent.rlm"] = rlm_package + +if "deadend_agent.tools" not in sys.modules: + tools_package = types.ModuleType("deadend_agent.tools") + tools_package.__path__ = [str(PACKAGE_ROOT / "tools")] + sys.modules["deadend_agent.tools"] = tools_package + +if "deadend_agent.tools.python_interpreter" not in sys.modules: + pyi_package = types.ModuleType("deadend_agent.tools.python_interpreter") + pyi_package.__path__ = [str(PACKAGE_ROOT / "tools" / "python_interpreter")] + sys.modules["deadend_agent.tools.python_interpreter"] = pyi_package + +if "deadend_agent.tools.python_interpreter.python_interpreter" not in sys.modules: + stub_python_interpreter = types.ModuleType("deadend_agent.tools.python_interpreter.python_interpreter") + + class PythonInterpreter: # pragma: no cover - import stub only + pass + + stub_python_interpreter.PythonInterpreter = PythonInterpreter + sys.modules["deadend_agent.tools.python_interpreter.python_interpreter"] = stub_python_interpreter + +if "deadend_agent.core_agent" not in sys.modules: + core_agent_package = types.ModuleType("deadend_agent.core_agent") + core_agent_package.__path__ = [str(PACKAGE_ROOT / "core_agent")] + + class LLMError(Exception): + def __init__(self, message: str, original_error: Exception | None = None): + self.message = message + self.original_error = original_error + super().__init__(message) + + class RateLimitError(LLMError): + pass + + class QuotaExceededError(LLMError): + pass + + class AuthenticationError(LLMError): + pass + + class ConnectionError(LLMError): + pass + + class ModelNotFoundError(LLMError): + pass + + class InvalidRequestError(LLMError): + pass + + core_agent_package.LLMError = LLMError + core_agent_package.RateLimitError = RateLimitError + core_agent_package.QuotaExceededError = QuotaExceededError + core_agent_package.AuthenticationError = AuthenticationError + core_agent_package.ConnectionError = ConnectionError + core_agent_package.ModelNotFoundError = ModelNotFoundError + core_agent_package.InvalidRequestError = InvalidRequestError + sys.modules["deadend_agent.core_agent"] = core_agent_package + +from deadend_agent.core_agent.rlm_runner import SandboxedRLMRunner + + +def _exec_script(script: str) -> None: + globals_dict = {"__name__": "__main__"} + exec(script, globals_dict, globals_dict) + + +def test_runner_extracts_python_blocks_and_direct_final(tmp_path): + runner = SandboxedRLMRunner(root_model="openai/test", workspace_root=tmp_path) + blocks = runner._extract_code_blocks( + "before\n```python\nx = 1\n```\nafter\n```repl\nobserve(x)\n```" + ) + + assert blocks == ["x = 1", "observe(x)"] + assert runner._parse_direct_final("FINAL('done')") == "done" + + +def test_runner_script_persists_state_and_reuses_subcall_results(tmp_path): + runner = SandboxedRLMRunner(root_model="openai/test", workspace_root=tmp_path, session_id="session") + runner._prepare_workspace(memory_root=None, context={"docs": ["alpha", "beta"]}) + + first_script = runner._build_script( + "x = 7\n" + "observe('ready', x)\n" + "token = llm_query('summarize', 'hello world')\n" + "observe(token)\n" + ) + _exec_script(first_script) + + observation = json.loads((runner.artifacts_dir / "observation.json").read_text(encoding="utf-8")) + pending = json.loads((runner.artifacts_dir / "pending_subcalls.json").read_text(encoding="utf-8")) + + assert observation["observations"][0] == "ready 7" + assert observation["state_keys"] == ["token", "x"] + assert len(pending) == 1 + + request_id = pending[0]["request_id"] + (runner.artifacts_dir / "subcall_results.json").write_text( + json.dumps( + { + request_id: { + "prompt": "summarize", + "content": "hello world", + "model": "openai/test", + "result": "summary-result", + } + }, + ensure_ascii=False, + indent=2, + ), + encoding="utf-8", + ) + + second_script = runner._build_script( + "observe('state', x)\n" + "answer = llm_query('summarize', 'hello world')\n" + "observe(answer)\n" + "FINAL_VAR('answer')\n" + ) + _exec_script(second_script) + + final_payload = json.loads((runner.artifacts_dir / "final.json").read_text(encoding="utf-8")) + observation = json.loads((runner.artifacts_dir / "observation.json").read_text(encoding="utf-8")) + + assert final_payload["value"] == "summary-result" + assert observation["observations"][-1] == "summary-result" + + +def test_runner_script_supports_mount_based_file_ops(tmp_path): + runner = SandboxedRLMRunner(root_model="openai/test", workspace_root=tmp_path, session_id="session-avfs") + runner._prepare_workspace(memory_root=None, context={"seed": "ok"}) + + script = runner._build_script( + "mounted = mount('.')\n" + "observe('mounted', mounted)\n" + "write_file('notes/a.txt', 'alpha\\nbeta\\n')\n" + "files = list_files('notes')\n" + "observe('files', files)\n" + "snippet = read_lines('notes/a.txt', 2, 2)\n" + "observe('line2', snippet)\n" + "hits = grep('alpha', path='notes')\n" + "observe('hits', hits[0]['path'], hits[0]['line_number'])\n" + ) + _exec_script(script) + + observation = json.loads((runner.artifacts_dir / "observation.json").read_text(encoding="utf-8")) + observed = "\n".join(observation["observations"]) + + assert "mounted " in observed + assert "notes/a.txt" in observed + assert "line2 beta" in observed + assert "hits notes/a.txt 1" in observed diff --git a/deadend_cli/deadend_eval/src/deadend_eval/eval.py b/deadend_cli/deadend_eval/src/deadend_eval/eval.py index a330b4a..d535504 100644 --- a/deadend_cli/deadend_eval/src/deadend_eval/eval.py +++ b/deadend_cli/deadend_eval/src/deadend_eval/eval.py @@ -20,8 +20,9 @@ Sandbox, DeadEndAgent, ) -from deadend_agent.config.settings import ModelSpec +from deadend_agent.config.settings import Config, ModelSpec from deadend_agent.models.registry import EmbedderClient +from deadend_agent.utils.network import deterministic_session_id from deadend_eval.metrics import ( instrument_agent_runner, global_metrics, @@ -166,11 +167,16 @@ async def eval_deadend_agent( target_host = eval_metadata.target_host session_id = uuid4() + local_agent_id = Config.get_local_agent_id() + embedding_session_id = deterministic_session_id(eval_metadata.target_host or "localhost") deadend_agent = DeadEndAgent( session_id=session_id, + embedding_session_id=embedding_session_id, model=model, available_agents=generic_agents, - max_depth=2 + max_depth=2, + agents_storage_root=Config.agents_storage_root, + local_agent_id=local_agent_id, ) # Set challenge name for trace file naming deadend_agent.challenge_name = eval_metadata.name diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 new file mode 100644 index 0000000..f02704a --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 @@ -0,0 +1,21 @@ +## MEMORY SUMMARY + +At the end of your work, append a concise summary to the persistent memory workspace using `avfs_write(..., workspace="memory")`. +For the memory workspace, you may only use `avfs_write`. Do not use `avfs_list`, `avfs_read`, or `avfs_grep` against `workspace="memory"`. + +Requirements: +- Write only factual summaries derived from the work you actually performed. +- Include what you tested, what worked, what failed, and the best next step. +- Append to a stable file for your agent type so later runs can build on it. + +Recommended path: +- `summaries/{{agent_name}}.md` + +Recommended format: +```text +## Task Summary +- Objective: ... +- Actions: ... +- Findings: ... +- Next step: ... +``` diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 new file mode 100644 index 0000000..cf16fcc --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 @@ -0,0 +1,179 @@ +You are a memory management agent. Your only job is to inspect and update the persistent agent memory workspace. + +Use the available tools to: +- list existing notes and summaries +- read prior summaries before writing new notes +- append concise, factual summaries +- keep files organized and readable + +Do not perform network requests, shell execution, or target testing. Only operate on the memory workspace. + +When you call AVFS tools, always target the memory workspace explicitly with `workspace="memory"`. +Your main startup task is to inspect prior memory and produce a concise context that other agents can use immediately. + +## RETRIEVAL PRIORITIES + +Your retrieval strategy must depend on the current task or user query. +Do not summarize the whole memory workspace indiscriminately. Retrieve the most important information for the task at hand. + +## TOOL ORDERING + +Use the tools in a deliberate order based on the query: + +1. Extract search patterns from the current task or user query +- Pull out the highest-signal terms: target names, endpoints, parameters, bug classes, payload keywords, files, credentials, technologies, tool names +- Keep both exact phrases and a few compact fallback patterns + +2. Use `avfs_grep(..., workspace="memory")` first +- Search the memory workspace for the patterns you extracted +- Use grep to narrow the candidate files before reading +- Prefer targeted grep over broad file-by-file reading + +3. Use `avfs_read(..., workspace="memory")` on the most relevant files +- Read only the files that appear strongly connected to the task +- Start with files that contain exact matches, repeated matches, or high-signal evidence +- Expand reading only if the first pass is insufficient + +4. Use `avfs_list(..., workspace="memory")` when structure matters +- Use listing to understand organization, discover likely summary files, or confirm emptiness +- Do not list recursively without a reason; use it to orient the search, not replace grep + +5. Build relations across the retrieved evidence +- Connect files, findings, payloads, endpoints, credentials, constraints, and outcomes +- Identify which items reinforce each other, contradict each other, or form an execution chain +- Surface the most useful relationships for the current task, not every possible relationship + +6. Return a compact, task-directed context +- Summarize the evidence chain, not just isolated notes +- Explain what the relationship between the retrieved items implies for the next step + +Prioritize memory in this order: +1. Information directly related to the current objective, target, endpoint, exploit path, bug class, or requested action +2. Previously successful actions, working payloads, confirmed findings, and concrete evidence +3. Known failures, dead ends, blocked approaches, and constraints that should prevent wasted retries +4. Reusable context that can accelerate execution: credentials, session details, relevant files, tool outputs, environment assumptions +5. Lower-signal background notes only if they materially help with the current task + +When selecting what to return: +- Prefer precise, actionable facts over broad summaries +- Prefer recent or repeated signals over one-off speculation +- Prefer confirmed evidence over tentative hypotheses +- Exclude stale or weakly related notes unless they change execution decisions + +If multiple memory items compete, rank them by: +- relevance to the current task +- confirmed usefulness +- specificity +- actionability for downstream agents + +If the user query is narrow, your output should be narrow. +If the user query is broad, return the smallest set of memory needed to move the task forward. + +Your output should help downstream agents answer: +- What already worked? +- What already failed? +- What should be reused? +- What should be avoided? +- What is the fastest credible next move? + +## RELATION BUILDING + +Do not stop at file retrieval. After grep and read: +- connect matching endpoints to successful or failed payloads +- connect credentials or session artifacts to the actions they enabled +- connect repeated findings across files into stronger evidence +- connect contradictions and uncertainty to concrete cautions + +Prefer statements like: +- "This payload worked on this endpoint according to these files" +- "These notes contradict each other, so the result is uncertain" +- "This credential was used successfully in the same exploit path" + +Avoid disconnected bullet dumps when a relationship can be made explicit. + +## EXAMPLES + +Example: retrieving prior vulnerability work +- Query intent: "Look for previous SQL injection findings on the login flow" +- Extract likely patterns: `sqli`, `sql injection`, `/login`, `login`, `auth`, `bypass` +- First action: grep those patterns across the memory workspace +- Then read only the files with strong matches +- Return: + - the endpoints involved + - the payloads that worked or failed + - the evidence proving the finding + - what should be retried or avoided + +Example: retrieving endpoint-specific recollection +- Query intent: "What do we already know about /api/users?" +- Extract likely patterns: `/api/users`, `users`, `idor`, `auth`, `token` +- Grep first, then read the files that mention that endpoint repeatedly +- Build relations between: + - the endpoint + - required authentication + - parameters + - previously observed responses + - vulnerabilities suspected or confirmed + +Example: retrieving exploit-path memory +- Query intent: "Find anything relevant to stored XSS in comments" +- Extract likely patterns: `stored xss`, `xss`, `comment`, `comments`, `payload`, `admin view` +- Grep for the vulnerability class, the feature name, and likely trigger locations +- Read the most relevant files +- Return the exploit chain if present: + - injection point + - payload used + - where execution happened + - what evidence confirmed execution + +Example: retrieving reusable credentials or sessions +- Query intent: "Do we have credentials or session artifacts for this target?" +- Extract likely patterns: `credential`, `password`, `cookie`, `session`, `token`, `jwt`, target name, endpoint names +- Grep first to locate candidate files +- Read only files with direct references to authentication material +- Return only concrete reusable details and the actions they enabled + +Example: handling weak or empty memory +- Query intent: "What do we know about SSTI on this target?" +- Grep returns no meaningful matches, and listing shows little or no stored memory +- Return that memory is empty or insufficient for SSTI recollection +- Do not invent prior tests +- If useful, suggest that downstream agents start from fresh reconnaissance + +## EMPTY MEMORY HANDLING + +Start by checking the memory workspace with `avfs_list(..., workspace="memory")`. + +If the memory workspace is empty, treat that as a real signal: +- No prior action was saved for this agent/session pair, or +- No reliable recollection can be reconstructed from persistent memory + +In that case: +- Do not infer history +- Do not invent previous findings, endpoints, credentials, or next steps +- Return a clearly empty memory state so downstream agents know they are starting fresh + +Recommended empty response behavior: +- Return a short plain-text statement that the memory workspace is empty and there is no prior persisted context to reuse +- If useful, add one short sentence saying the task should start from fresh reconnaissance + +If the workspace contains files but they are too sparse, stale, contradictory, or low-signal to support a reliable recollection: +- Say that the memory is insufficient or inconclusive +- Prefer partial, explicitly qualified context over confident synthesis +- Keep uncertainty visible instead of smoothing it away + +Your job is not to sound complete. Your job is to be trustworthy for downstream agents. + +## AVAILABLE TOOLS + +{% for tool_name, tool_description in tools.items() %} +### {{tool_name}} +{{tool_description}} +{% endfor %} + +## OUTPUT REQUIREMENTS + +- Return plain text only +- The text should be a concise task-relevant memory summary for the supervisor +- If memory is empty or irrelevant, say so explicitly in plain text +- Mention the most important prior successes, failures, reusable artifacts, and cautions only if they are relevant to the current task diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 index 7eba045..9207b33 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 @@ -176,15 +176,14 @@ This is an isolated container with pentesting tools - NOT a general-purpose shel - Do **NOT** refuse or avoid using a tool (e.g. `curl`, `nmap`, `sqlmap`, `hydra`, `ffuf`, etc.) just because it is powerful or could be dangerous on a non-sandboxed system; here it is safe by design. - Your only safety filters are **scope** (what targets are allowed) and the explicit constraints in these instructions. -## CRITICAL CONSTRAINT - NO SOURCE CODE ACCESS +## CRITICAL CONSTRAINT - NO RAW SOURCE CODE ACCESS THROUGH SHELL -**This container has NO access to target source code files.** +**Do not use the shell to inspect source or workspace files.** -- **DO NOT** attempt to read, grep, cat, or find source code files -- **DO NOT** look for .py, .js, .html, .php, or any code files in this container -- **DO NOT** use file system commands to explore source code - -If your task requires source code analysis, return with low confidence and note that source code tools are needed. +- **DO NOT** use `cat`, `grep`, `find`, `sed`, `awk`, or similar shell file commands on source/workspace files +- **DO NOT** inspect `.py`, `.js`, `.html`, `.php`, or other code files through raw shell commands +- **DO NOT** use shell filesystem exploration for source browsing +- If a task requires source/workspace inspection, return with low confidence and state that shell is not the right agent for that task ## AVAILABLE TOOLS diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/template_renderer.py b/deadend_cli/deadend_prompts/src/deadend_prompts/template_renderer.py index 73576af..f7be27b 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/template_renderer.py +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/template_renderer.py @@ -35,7 +35,7 @@ def get_instructions(self, **kwargs): _metadata, instructions = _parse_template_metadata(instructions) instructions_template = self.env.from_string(instructions) - return instructions_template.render(tools=self.tools, **kwargs) + return instructions_template.render(agent_name=self.agent_name, tools=self.tools, **kwargs) def get_preprompt(self, **kwargs): raise NotImplementedError diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_grep.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_grep.description.jinja2 new file mode 100644 index 0000000..7054b96 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_grep.description.jinja2 @@ -0,0 +1 @@ +Search the mounted AVFS workspace only using ripgrep via ripgrepy. Use this for safe, efficient code and text search inside the project/workspace, not for system paths outside the mounted workspace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_list.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_list.description.jinja2 new file mode 100644 index 0000000..0c6294c --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_list.description.jinja2 @@ -0,0 +1 @@ +List files and directories inside the mounted AVFS workspace only. Use this for mounted project/workspace contents, not for system paths like `/usr`, `/etc`, or other non-workspace locations. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_read.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_read.description.jinja2 new file mode 100644 index 0000000..4baa2a4 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_read.description.jinja2 @@ -0,0 +1 @@ +Read a text file from the mounted AVFS workspace only, with optional line bounds. Use this for project/workspace inspection instead of shell commands like `cat` or `sed`, not for system paths outside the mounted workspace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_write.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_write.description.jinja2 new file mode 100644 index 0000000..b57f6d0 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/avfs_write.description.jinja2 @@ -0,0 +1 @@ +Write or append text to a file inside the mounted AVFS workspace. Changes apply to the underlying host workspace rooted in the configured AVFS workspace. diff --git a/deadend_cli/pyproject.toml b/deadend_cli/pyproject.toml index 11029c9..553ec0b 100644 --- a/deadend_cli/pyproject.toml +++ b/deadend_cli/pyproject.toml @@ -57,6 +57,7 @@ dependencies = [ "arize-phoenix-evals>=2.10.0", "openinference-instrumentation>=0.1.44", "openinference-instrumentation-litellm>=0.1.29", + "ripgrepy>=2.2.0", ] [project.scripts] @@ -149,11 +150,13 @@ extraPaths = [ members = [ "deadend_agent", "deadend_prompts", - "deadend_eval" + "deadend_eval", + "simple-python-interpreter-sandbox", ] [tool.uv.sources] deadend-agent = { workspace = true } deadend-prompts = { workspace = true } deadend-eval = { workspace = true } +python-sandbox-client = { workspace = true } diff --git a/deadend_cli/simple-python-interpreter-sandbox b/deadend_cli/simple-python-interpreter-sandbox new file mode 160000 index 0000000..6c99fe5 --- /dev/null +++ b/deadend_cli/simple-python-interpreter-sandbox @@ -0,0 +1 @@ +Subproject commit 6c99fe52c39ad40d469c81af112a5cc689b49b9b diff --git a/deadend_cli/src/deadend_cli/chat.py b/deadend_cli/src/deadend_cli/chat.py index c35bf8b..2667ba1 100644 --- a/deadend_cli/src/deadend_cli/chat.py +++ b/deadend_cli/src/deadend_cli/chat.py @@ -356,6 +356,8 @@ async def chat_interface( # OpenAPI spec if available knowledge_base: str, # Knowledge base path + workspace_root: str | None = None, + # Optional AVFS workspace root llm_provider: str = "openai" # LLM provider ): @@ -389,6 +391,7 @@ async def chat_interface( 'requester': "Agent specialized in fine-grained testing and sending raw request data. Capable of handling authentication (session and token). Uses pupeteer in the background. Capable of exploring APIs and websites. Best for gathering auth tokens, testing individual endpoints, and precise request manipulation. Should NOT be used for automation tasks such as fuzzing or repetitive tasks that need iteration - use python_interpreter for those tasks instead.", 'python_interpreter': "Agent specialized in generating code and running it. Each code generated is ran safely in a sandboxed webassembly. Best for fuzzing, parameter testing, generating testing exploits, and repetitive security testing operations that require iteration. Use this agent for tasks that need automation, loops, or multiple iterations.", 'shell': "Agent that gives access to a terminal bash shell. Run linux commands here.", + 'memory': "Agent specialized in reading and writing the persistent memory workspace under the agent cache.", 'router_agent': 'Router agent, expert that routes to the specific agent needed to achieve the next step of the plan.' } @@ -446,7 +449,7 @@ async def approval_callback(): console_printer.print("[red]No target provided. Exiting application...[/red]") sys.exit(1) - # Create agent with deterministic session ID per target + # Create agent with a runtime session plus deterministic target embedding session runtime_session_id = uuid4() embedding_session_id = deterministic_session_id(target) deadend_agent = DeadEndAgent( @@ -454,7 +457,10 @@ async def approval_callback(): embedding_session_id=embedding_session_id, model=model, available_agents=available_agents, - max_depth=3 + max_depth=3, + workspace_root=workspace_root, + agents_storage_root=config.agents_storage_root, + local_agent_id=local_agent_id, ) deadend_agent.set_approval_callback(approval_callback) @@ -570,7 +576,10 @@ def interrupt_agent(): embedding_session_id=embedding_session_id, model=model, available_agents=available_agents, - max_depth=3 + max_depth=3, + workspace_root=workspace_root, + agents_storage_root=config.agents_storage_root, + local_agent_id=local_agent_id, ) deadend_agent.set_approval_callback(approval_callback) # Re-initialize with new target diff --git a/deadend_cli/src/deadend_cli/cli.py b/deadend_cli/src/deadend_cli/cli.py index 05c3e87..be1e021 100644 --- a/deadend_cli/src/deadend_cli/cli.py +++ b/deadend_cli/src/deadend_cli/cli.py @@ -60,6 +60,7 @@ def chat( None, help="Path to the OpenAPI specification file." ), knowledge_base: str = typer.Option(None, help="Folder path to the knowledge base."), + workspace_root: str = typer.Option(None, help="Host workspace to mount into AVFS."), ): """Run the interactive chat agent. @@ -106,6 +107,7 @@ def chat( target=target, openapi_spec=openapi_spec, knowledge_base=knowledge_base, + workspace_root=workspace_root, ) ) finally: diff --git a/deadend_cli/src/deadend_cli/jsonrpc_server.py b/deadend_cli/src/deadend_cli/jsonrpc_server.py index 668fc77..596715b 100644 --- a/deadend_cli/src/deadend_cli/jsonrpc_server.py +++ b/deadend_cli/src/deadend_cli/jsonrpc_server.py @@ -502,6 +502,7 @@ async def instantiate_agent( # Get provider and model from params, or use current defaults provider = params.get("provider") model_name = params.get("model_name") + workspace_root = params.get("workspace_root") # Get the model spec (will use current provider/model if not specified) logger.info("model and provider %s %s", provider, model_name) @@ -528,6 +529,7 @@ async def instantiate_agent( "Best for fuzzing, parameter testing, and repetitive security testing operations." ), "shell": "Agent providing access to a bash shell for running Linux commands.", + "memory": "Agent specialized in reading and writing the persistent memory workspace under the agent cache.", "router_agent": "Router agent that selects the appropriate specialized agent.", "webapp_analyzer": ( "Front-end webapp analyzer. This agent is specialized in looking into the web application" @@ -540,7 +542,10 @@ async def instantiate_agent( embedding_session_id=embedding_session_id, model=model, available_agents=available_agents, - max_depth=3 + max_depth=3, + workspace_root=workspace_root, + agents_storage_root=component_manager.config.agents_storage_root, + local_agent_id=component_manager.config.get_local_agent_id(), ) async def approval_callback() -> str: return "yes" diff --git a/deadend_cli/uv.lock b/deadend_cli/uv.lock index 0c2cca7..928be00 100644 --- a/deadend_cli/uv.lock +++ b/deadend_cli/uv.lock @@ -17,6 +17,7 @@ members = [ "deadend-cli", "deadend-eval", "deadend-prompts", + "python-sandbox-client", ] [[package]] @@ -979,9 +980,11 @@ dependencies = [ { name = "pydantic" }, { name = "pydantic-ai" }, { name = "pydantic-ai-slim", extra = ["google", "openrouter"] }, + { name = "python-sandbox-client" }, { name = "pyyaml" }, { name = "readchar" }, { name = "rich" }, + { name = "ripgrepy" }, { name = "semantic-text-splitter" }, { name = "sqlalchemy" }, { name = "tenacity" }, @@ -1041,9 +1044,11 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.11.5" }, { name = "pydantic-ai", specifier = ">=1.35.0" }, { name = "pydantic-ai-slim", extras = ["google", "openrouter"], specifier = ">=1.35.0" }, + { name = "python-sandbox-client", editable = "simple-python-interpreter-sandbox" }, { name = "pyyaml", specifier = ">=6.0.2" }, { name = "readchar", specifier = ">=4.2.1" }, { name = "rich", specifier = ">=14.0.0" }, + { name = "ripgrepy", specifier = ">=2.2.0" }, { name = "semantic-text-splitter", specifier = ">=0.27.0" }, { name = "sqlalchemy", specifier = ">=2.0.41" }, { name = "tenacity", specifier = ">=9.1.2" }, @@ -1116,6 +1121,7 @@ dependencies = [ { name = "pyyaml" }, { name = "readchar" }, { name = "rich" }, + { name = "ripgrepy" }, { name = "semantic-text-splitter" }, { name = "sqlalchemy" }, { name = "tenacity" }, @@ -1185,6 +1191,7 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0.2" }, { name = "readchar", specifier = ">=4.2.1" }, { name = "rich", specifier = ">=14.0.0" }, + { name = "ripgrepy", specifier = ">=2.2.0" }, { name = "semantic-text-splitter", specifier = ">=0.27.0" }, { name = "sqlalchemy", specifier = ">=2.0.41" }, { name = "tenacity", specifier = ">=9.1.2" }, @@ -4056,6 +4063,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/45/58/38b5afbc1a800eeea951b9285d3912613f2603bdf897a4ab0f4bd7f405fc/python_multipart-0.0.20-py3-none-any.whl", hash = "sha256:8a62d3a8335e06589fe01f2a3e178cdcc632f3fbe0d492ad9ee0ec35aab1f104", size = 24546, upload-time = "2024-12-16T19:45:44.423Z" }, ] +[[package]] +name = "python-sandbox-client" +version = "0.1.0" +source = { editable = "simple-python-interpreter-sandbox" } + [[package]] name = "pytokens" version = "0.2.0" @@ -4340,6 +4352,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/13/2f/b4530fbf948867702d0a3f27de4a6aab1d156f406d72852ab902c4d04de9/rich_rst-1.3.2-py3-none-any.whl", hash = "sha256:a99b4907cbe118cf9d18b0b44de272efa61f15117c61e39ebdc431baf5df722a", size = 12567, upload-time = "2025-10-14T16:49:42.953Z" }, ] +[[package]] +name = "ripgrepy" +version = "2.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/80/8a/023e7c432634a6090a26ace319a37a2a73aac8fa6a7bf142fd7b3ec8fd3b/ripgrepy-2.2.0.tar.gz", hash = "sha256:4c43c61384f257660007acd271a5d8e4abe9be0b069c418d091f7299e080ca9d", size = 31740, upload-time = "2025-07-11T01:18:01.662Z" } + [[package]] name = "rpds-py" version = "0.27.1" diff --git a/docs/RLM.md b/docs/RLM.md new file mode 100644 index 0000000..570f5e9 --- /dev/null +++ b/docs/RLM.md @@ -0,0 +1,359 @@ +# Step-by-step implementation guide from **Recursive Language Models** + +## Resource List +- **Recursive Language Models** — Alex L. Zhang, Tim Kraska, Omar Khattab, **2026** (arXiv v2, revised January 28, 2026), **arXiv:2512.24601**. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +## Title: Recursive Language Models + +**Citation:** Zhang, Kraska, Khattab, 2026, arXiv:2512.24601 (v2). ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +**Summary:** +The paper proposes a long-context inference scaffold called a **Recursive Language Model (RLM)**. Instead of feeding the full long prompt into the model, the prompt is stored in an external environment, and the LM interacts with it through a persistent Python REPL. The LM can inspect the context, transform it with code, and call a sub-LM on selected chunks. This lets the system handle inputs far beyond the base model’s context window while preserving more fine-grained access than summarization-based methods. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +### Research Question / Problem +The paper asks how to let LLMs process **arbitrarily long prompts** without relying only on larger context windows or lossy summarization. The authors argue that long-context failure depends not just on input length but also on task complexity, and that many tasks need dense access to the prompt rather than a compressed summary. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +### Method & Reasoning +The core design is simple: + +1. Treat the long prompt as part of the **external environment**, not as direct LM input. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +2. Start a **persistent Python REPL** with a `context` variable containing the prompt and a helper like `llm_query(...)` for sub-LM calls. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +3. Let the root LM iteratively emit REPL code, inspect outputs, maintain variables/buffers, and eventually return an answer using `FINAL(...)` or `FINAL_VAR(...)`. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +4. Use sub-LM calls when semantic interpretation is needed on chunks that are still too large or too dense for the root LM to reason about directly. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +The paper’s reason for this architecture is that code execution gives the LM a symbolic way to search, filter, partition, and aggregate over huge inputs, while recursive LM calls let it delegate hard semantic work on manageable subproblems. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +### Findings & Chain of Evidence +Empirically, the RLM scaffold outperformed direct model calls and several long-context baselines across CodeQA, BrowseComp+, OOLONG, and OOLONG-Pairs. On the GPT-5 setup, the RLM beat the base model on all four listed benchmarks and handled BrowseComp+ inputs in the 6M–11M token range that the base model could not fit directly. The paper also shows that the REPL alone already helps with long inputs, while recursive sub-calls are especially important for **information-dense** tasks like OOLONG and OOLONG-Pairs. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +The qualitative trajectory analysis supports the implementation choices: successful runs often use regex or lightweight probing first, then chunking, then sub-LM calls over selected pieces, then programmatic aggregation into buffers or final variables. The paper also shows failure cases where models over-verify, make too many sub-calls, or fail to return a prepared variable cleanly. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +### Strengths & Limitations +**Strengths** +- Handles contexts far beyond the base context window. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Preserves fine-grained access to source material better than iterative summarization. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Supports long outputs by building them in variables and returning `FINAL_VAR(...)`. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Works as an inference scaffold without retraining the base model. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +**Limitations** +- The paper’s implementation uses **blocking/sequential** sub-calls, which hurts runtime. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Final-answer detection via `FINAL(...)` / `FINAL_VAR(...)` is brittle. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Models with weak coding ability or insufficient output budget perform poorly as RLMs. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- The paper only uses **max recursion depth 1** in experiments, so “recursive” here is operationally shallow in the reported system. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +--- + +## Integration Guidelines + +If you want to implement the paper faithfully, the safest interpretation is: + +- Build an **agent scaffold**, not a new neural architecture. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Use a **persistent Python execution environment** as the working memory. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Expose at minimum: + - `context` + - `print()` + - `llm_query(prompt)` for sub-LM inference. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +- Prompt the root LM to: + - inspect context first, + - choose a chunking/filtering strategy, + - batch work into sub-calls, + - store intermediate results in variables, + - return only when ready with `FINAL(...)` or `FINAL_VAR(...)`. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +The paper also suggests two practical heuristics: +1. For sparse-retrieval-style tasks, start with **regex/keyword probing** and only sub-call on likely relevant snippets. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) +2. For information-dense tasks, do **systematic chunking plus sub-LM semantic labeling**, then aggregate programmatically. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +--- + +# Step-by-step implementation blueprint + +## 1. Define the external interface +Your RLM should look like a normal function: + +```python +answer = rlm(query: str, context: Any) -> str +``` + +The paper says the RLM should preserve the same external abstraction as an LM: accept a prompt/context and return a string answer, while internally offloading the context into the environment. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## 2. Load the context into a persistent REPL +Create a sandboxed Python session with persistent variables: + +```python +state = { + "context": context, +} +``` + +Also compute metadata the paper includes in the system prompt: +- `context_type` +- `context_total_length` +- `context_lengths` for chunks/documents/sections. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +A practical implementation: + +```python +def describe_context(context): + if isinstance(context, str): + return { + "context_type": "string", + "context_total_length": len(context), + "context_lengths": [len(context)], + } + elif isinstance(context, list): + return { + "context_type": "List[str]", + "context_total_length": sum(len(x) for x in context), + "context_lengths": [len(x) for x in context], + } + elif isinstance(context, dict): + # adapt as needed + text = str(context) + return { + "context_type": "dict", + "context_total_length": len(text), + "context_lengths": [len(text)], + } +``` + +This mirrors the prompt fields shown in Appendix D. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## 3. Expose the REPL tools +The minimum environment from the paper is: + +- `context` +- `llm_query(...)` +- `print(...)` with visible truncated output. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +A minimal helper layer: + +```python +class RLMEnv: + def __init__(self, context, sub_llm, max_print_chars=4000): + self.globals = {"context": context} + self.sub_llm = sub_llm + self.max_print_chars = max_print_chars + + def llm_query(self, prompt: str) -> str: + return self.sub_llm(prompt) + + def run(self, code: str) -> str: + import io, contextlib + buf = io.StringIO() + self.globals["llm_query"] = self.llm_query + with contextlib.redirect_stdout(buf): + exec(code, self.globals, self.globals) + out = buf.getvalue() + return out[:self.max_print_chars] +``` + +The truncation behavior is important because the paper explicitly tells the model that REPL outputs are truncated and that variables should be used as buffers for larger intermediate state. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## 4. Use the paper’s root system prompt +The most implementation-critical part of the paper is the RLM system prompt. At minimum, keep these behaviors: + +- The LM is told it can analyze context in a REPL. +- It is told `context` contains key information. +- It is told `llm_query` can query a sub-LM with about **500K characters** of context. +- It is told to inspect the context, choose chunking, use buffers, and finish with `FINAL(...)` or `FINAL_VAR(...)`. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +A faithful condensed version: + +```text +You are tasked with answering a query with associated context. +You can access, transform, and analyze this context interactively in a REPL environment +that can recursively query sub-LLMs. You will be queried iteratively until you provide +a final answer. + +Your context is a {context_type} with {context_total_length} total characters, +and is broken up into chunks of char lengths: {context_lengths}. + +The REPL environment is initialized with: +1. A `context` variable that contains extremely important information. +2. A `llm_query` function that allows you to query an LLM inside your REPL environment. +3. The ability to use `print()` statements to view REPL output. + +You will only be able to see truncated outputs from the REPL environment. +Use variables as buffers to build up your final answer. + +Make sure to explicitly look through the entire context in REPL before answering your query. +A good strategy is to inspect the context, choose a chunking strategy, query an LLM per chunk, +store answers in buffers, and aggregate them. + +When you want to execute Python code, wrap it in ```repl ... ```. + +When done, return either: +FINAL(your final answer here) +or +FINAL_VAR(variable_name) +``` + +That is a paraphrased implementation-oriented reduction of the Appendix D prompt. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## 5. Build the iterative root loop +The root LM should operate in turns: + +1. Send system prompt + user query. +2. Get assistant output. +3. If it emits `repl` code blocks, execute them in the persistent environment. +4. Return the execution output as the next user/tool message. +5. Repeat until `FINAL(...)` or `FINAL_VAR(...)` appears. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +Pseudo-logic: + +```python +def run_rlm(root_llm, sub_llm, query, context, max_turns=30): + meta = describe_context(context) + env = RLMEnv(context, sub_llm) + messages = [ + {"role": "system", "content": build_system_prompt(meta)}, + {"role": "user", "content": query}, + ] + + for _ in range(max_turns): + assistant = root_llm(messages) + text = assistant["content"] + + final = parse_final(text) + if final is not None: + return final.resolve(env.globals) + + code_blocks = extract_repl_blocks(text) + if code_blocks: + outputs = [] + for code in code_blocks: + outputs.append(env.run(code)) + messages.append({"role": "assistant", "content": text}) + messages.append({"role": "user", "content": "\n\n".join(outputs)}) + else: + messages.append({"role": "assistant", "content": text}) + messages.append({ + "role": "user", + "content": "Continue reasoning in REPL, or return FINAL(...) / FINAL_VAR(...)." + }) + + raise RuntimeError("RLM did not terminate") +``` + +The explicit final tags are paper-faithful, but the authors also report that this termination method is brittle, so add parser safeguards and a max-turn budget. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## 6. Implement `FINAL(...)` and `FINAL_VAR(...)` +The paper’s prompt requires two completion modes: + +- `FINAL(answer text)` +- `FINAL_VAR(variable_name)` to return a variable built inside the REPL. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +This is especially useful for long outputs, where the root model should build the answer incrementally in Python lists/strings and then return the variable instead of regenerating it from scratch. That pattern is directly motivated by the OOLONG-Pairs behavior described in the paper. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## 7. Keep recursion shallow at first +For a faithful reproduction, make sub-calls plain LM calls rather than full nested RLM calls. The paper’s experiments used **max recursion depth 1**, meaning sub-calls were LMs, not deeper recursive environments. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +So: + +```python +def llm_query(prompt): + return sub_llm(prompt) # no nested REPL by default +``` + +That is the simplest way to match the paper before experimenting with deeper recursion. + +## 8. Choose root and sub models deliberately +The paper used stronger models with coding ability as root controllers and, in at least the GPT-5 setup, a cheaper smaller model for sub-calls: GPT-5 for the root and GPT-5-mini for recursive calls. It also reports that smaller models with weak coding ability struggled as RLMs. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +Implementation takeaway: +- **Root LM:** prioritize coding ability, tool use, long reasoning. +- **Sub LM:** prioritize cost-efficiency and sufficient context length. + +## 9. Add batching constraints to `llm_query` +The Qwen-specific prompt tweak in the paper is operationally important: batch information aggressively and avoid many tiny sub-calls. The added guidance recommends aiming for roughly **200K characters per call** and warns against issuing one sub-call per line or item. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +So your helper should support bulk chunks: + +```python +def batch_strings(items, target_chars=200_000): + batch, batches, size = [], [], 0 + for item in items: + if size + len(item) > target_chars and batch: + batches.append(batch) + batch, size = [], 0 + batch.append(item) + size += len(item) + if batch: + batches.append(batch) + return batches +``` + +## 10. Encode the paper’s common successful strategies as defaults +From the trajectory analysis, three patterns should be turned into implementation priors: + +### A. Probe first, then narrow +Start with small inspections: +- print the first few lines, +- inspect structure, +- regex-search likely keywords, +- identify candidate chunks. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +### B. Use sub-LM calls for semantic transforms +If the task requires label inference, meaning extraction, or question answering over chunks, send the chunk plus a focused question to `llm_query`. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +### C. Aggregate programmatically +Have the root LM use Python to: +- collect chunk answers, +- count labels, +- deduplicate entities/pairs, +- build final formatted output. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +--- + +# Recommended MVP implementation plan + +## Phase 1: Minimal faithful reproduction +Implement: +- persistent Python REPL +- `context` +- `llm_query` +- iterative loop +- `FINAL` / `FINAL_VAR` +- output truncation +- max turns and max sub-call budget. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## Phase 2: Production hardening +Add: +- sandboxing for code execution, +- async sub-calls, +- per-query cost budget, +- retry/repair when code fails, +- anti-loop heuristics when repeated verification appears. The paper explicitly points to asynchronous sub-calls and sandboxed REPLs as promising improvements over its own implementation. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +## Phase 3: Routing logic +Because the paper notes that base LMs can outperform RLMs on smaller contexts, add a simple router: +- small/simple prompt → direct LM +- large or information-dense prompt → RLM. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +--- + +# Practical design choices the paper leaves underspecified + +Based on the text available in the paper, these items are **not fully specified**, so you will need to choose them yourself: + +- exact REPL sandbox implementation, +- exact stdout truncation length, +- exact parsing grammar for `FINAL(...)`, +- retry behavior for malformed code, +- message format between assistant turn and REPL feedback, +- whether sub-LM calls inherit any special system prompt. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +So if your goal is a working system, treat the paper as a **strong scaffold design**, not a drop-in full spec. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +--- + +# Cross-Resource Synthesis +For this paper, the main actionable insight is: **move long context out of the LM’s token stream and into an executable environment**. Then let the LM decide how to inspect, chunk, delegate, and aggregate. The paper’s experiments support three implementation priorities: use a persistent REPL, let the model call a sub-LM on selected chunks, and keep intermediate state in variables rather than repeatedly re-generating answers. The biggest engineering upgrades beyond the paper are likely async sub-calls, stronger loop control, and safer sandboxing. ([ar5iv.org](https://ar5iv.org/html/2512.24601v2)) + +If you want, I can next turn this into one of these: +1. a **Python skeleton implementation**, +2. an **OpenAI API-based version**, or +3. a **LangGraph / agent-framework version**. \ No newline at end of file diff --git a/docs/RLM_as_memory.md b/docs/RLM_as_memory.md new file mode 100644 index 0000000..af25b41 --- /dev/null +++ b/docs/RLM_as_memory.md @@ -0,0 +1,663 @@ +Below is a **practical implementation interpretation** of the paper for **agent memory over long Markdown (`.md`) and JSON (`.json`) files**. + +The key paper idea is to **not stuff the whole long context into the model prompt**. Instead, treat the long context as an **external environment** that the model can inspect, decompose, and query recursively over selected snippets. That is the right mental model for agent memory too: memory should be **stored externally, navigated programmatically, and only selectively exposed to the LM**. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# Step-by-step implementation for agent memory from long `.md` and `.json` files + +## 1. Core idea: use files as external memory, not prompt text + +For agent memory, the paper maps well to this pattern: + +- **Memory store** = files on disk / object store / DB +- **Working memory** = a Python environment with handles to those files +- **Reasoning loop** = root LM decides what to inspect +- **Semantic compression** = sub-LM calls on selected chunks +- **Final answer / plan** = built from buffers, not from re-reading everything + +This follows the paper’s main design: long context is treated as part of the environment, and the LM examines only relevant pieces rather than ingesting the whole thing at once. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +## 2. What “agent memory” means here + +For your use case, memory can be split into 3 practical types: + +### A. Documentation memory +Useful for `.md` files: +- README files +- design docs +- architecture notes +- meeting notes +- changelogs +- runbooks +- ADRs + +### B. Structured state memory +Useful for `.json` files: +- config files +- workflow state +- prior tool outputs +- event histories +- API responses +- cache entries +- task metadata + +### C. Episodic agent memory +Can be either `.md` or `.json`: +- previous conversations +- previous plans +- previous failures +- execution traces +- summaries of completed work + +The paper’s idea is especially good when this memory is **large, heterogeneous, and too long to fit cleanly in context**. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 3. Best usage patterns for `.md` and `.json` + +## Markdown memory is best for: +- human-written notes +- long-form explanations +- sectioned documents +- docs where headings matter +- retrieval by topic + +### Why +Markdown usually has natural structure: +- headings +- bullet lists +- code blocks +- sections +- subsections + +That makes it ideal for **structural chunking**: +- split by heading +- summarize per section +- recurse on relevant sections only + +This matches the paper’s examples where the LM first inspects context structure, then chooses a chunking strategy and queries chunks selectively. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +## JSON memory is best for: +- explicit state +- machine-readable observations +- logs +- repeated records +- nested objects +- tool outputs + +### Why +JSON is easier to: +- filter by key +- slice by path +- aggregate programmatically +- count, deduplicate, sort, compare + +This matches the paper’s strong result that the environment is valuable not just for reading, but for **programmatic manipulation** before asking the LM for semantic interpretation. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 4. Recommended memory architecture + +Use a **two-layer memory design**. + +## Layer 1: Raw memory store +Keep original files unchanged. + +Example: +```text +memory/ + docs/ + architecture.md + roadmap.md + incidents.md + state/ + workflow_state.json + tool_results.json + tasks.json + episodic/ + 2026-03-25-session-summary.md + prior_runs.json +``` + +## Layer 2: Derived navigation artifacts +Precompute lightweight metadata: + +- file manifest +- section index for markdown +- JSON path index +- embeddings per chunk +- keyword / BM25 index +- timestamps +- tags + +Example: +```json +{ + "file": "docs/architecture.md", + "type": "markdown", + "sections": [ + {"id": "sec1", "heading": "System Overview", "start": 0, "end": 1800}, + {"id": "sec2", "heading": "Memory Layer", "start": 1801, "end": 4200} + ] +} +``` + +This is important because the paper’s scaffold works best when the model can inspect **structure and chunk boundaries** before reading content in detail. That is directly aligned with the prompt fields describing context type and chunk lengths. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 5. Memory operations your agent should support + +Implement tools that let the LM inspect memory without reading it all. + +## Minimum tool set + +### For all files +- `list_files()` +- `get_file_metadata(path)` +- `grep_memory(pattern, path=None)` +- `read_chars(path, start, end)` +- `read_lines(path, start, end)` + +### For Markdown +- `list_md_sections(path)` +- `read_md_section(path, heading_or_id)` +- `read_md_outline(path)` +- `search_md_headings(query)` + +### For JSON +- `json_keys(path, json_path=None)` +- `json_get(path, json_path)` +- `json_search(path, key=None, value_contains=None)` +- `json_sample_array(path, json_path, start, end)` +- `json_schema(path)` + +### For semantic recursion +- `llm_query(prompt, content)` +- `summarize_chunk(chunk, task)` +- `classify_chunk(chunk, labels)` +- `extract_facts(chunk, schema)` + +This follows the paper’s pattern: the LM should be able to inspect, filter, decompose, and then call a sub-LM over selected snippets. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 6. Step-by-step implementation + +## Step 1: Create file loaders + +### Markdown loader +```python +from pathlib import Path + +def load_markdown(path: str) -> str: + return Path(path).read_text(encoding="utf-8") +``` + +### JSON loader +```python +import json +from pathlib import Path + +def load_json(path: str): + return json.loads(Path(path).read_text(encoding="utf-8")) +``` + +--- + +## Step 2: Build structure-aware chunkers + +## Markdown chunker +Prefer splitting by headings, not fixed token windows. + +```python +import re + +def split_markdown_sections(text: str): + pattern = r'^(#{1,6})\s+(.*)$' + lines = text.splitlines() + sections = [] + current = {"heading": "ROOT", "level": 0, "content": []} + + for line in lines: + m = re.match(pattern, line) + if m: + if current["content"]: + sections.append({ + "heading": current["heading"], + "level": current["level"], + "content": "\n".join(current["content"]).strip() + }) + current = { + "heading": m.group(2).strip(), + "level": len(m.group(1)), + "content": [] + } + else: + current["content"].append(line) + + if current["content"]: + sections.append({ + "heading": current["heading"], + "level": current["level"], + "content": "\n".join(current["content"]).strip() + }) + return sections +``` + +### Why this matters +For markdown, **semantic sections are the natural memory unit**. The paper explicitly suggests inspecting structure and choosing a chunking strategy rather than blindly feeding raw text. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +## Step 3: Build JSON path-based slicing + +```python +def json_keys(obj): + if isinstance(obj, dict): + return list(obj.keys()) + if isinstance(obj, list): + return [f"[{i}]" for i in range(min(len(obj), 20))] + return [] +``` + +Use JSONPath-like access: + +```python +def json_get(obj, path: str): + cur = obj + for part in path.strip(".").split("."): + if "[" in part and "]" in part: + key, idx = part[:-1].split("[") + if key: + cur = cur[key] + cur = cur[int(idx)] + else: + cur = cur[part] + return cur +``` + +### Why this matters +With JSON, the best pattern is: +1. inspect schema +2. identify relevant branches +3. read samples +4. aggregate in code +5. use LM only where semantics are needed + +That is very close to the paper’s successful trajectories: use code first, LM second. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +## Step 4: Build derived indexes + +## Markdown index +Store: +- headings +- section text length +- keywords +- embedding +- timestamps if present + +## JSON index +Store: +- top-level keys +- nested paths +- arrays and lengths +- event types +- record IDs +- embedding of flattened records where useful + +Example manifest: +```python +memory_manifest = { + "docs/architecture.md": { + "type": "markdown", + "sections": ["System Overview", "Memory Layer", "Failure Modes"] + }, + "state/tasks.json": { + "type": "json", + "top_keys": ["tasks", "updated_at", "owner"] + } +} +``` + +--- + +## Step 5: Create a persistent memory REPL + +The paper’s central scaffold uses a persistent REPL environment where the LM can store variables, inspect context, and call sub-LMs. For agent memory, your REPL should expose memory tools instead of one raw giant string. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +### Suggested environment +```python +class MemoryEnv: + def __init__(self, root_dir, sub_llm): + self.root_dir = Path(root_dir) + self.vars = {} + self.sub_llm = sub_llm + + def list_files(self): + return [str(p.relative_to(self.root_dir)) + for p in self.root_dir.rglob("*") if p.is_file()] + + def read_text(self, path): + return (self.root_dir / path).read_text(encoding="utf-8") + + def load_json(self, path): + import json + return json.loads((self.root_dir / path).read_text(encoding="utf-8")) + + def llm_query(self, instruction, content): + prompt = f"{instruction}\n\nCONTENT:\n{content}" + return self.sub_llm(prompt) +``` + +--- + +## Step 6: Give the root LM the right policy + +Your root LM should be instructed to do this: + +1. **Inspect available files first** +2. **Infer structure before reading deeply** +3. **Use code/tools to narrow search** +4. **Batch related content into sub-LM calls** +5. **Store summaries/facts in variables** +6. **Aggregate programmatically** +7. **Only then produce final answer** + +That is the same high-level behavior encouraged by the paper’s REPL prompt. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 7. Specific workflows + +## Workflow A: long Markdown memory + +### Example use case +“Find the latest architectural decisions and unresolved risks across all docs.” + +### Step-by-step +1. `list_files()` +2. filter `.md` +3. `read_md_outline()` or parse headings +4. find candidate sections: + - “Decision” + - “Architecture” + - “Risk” + - “Open Questions” +5. read only those sections +6. batch 3–10 sections into sub-LM summaries +7. store extracted facts: + - decision + - rationale + - risk + - owner + - status +8. merge and deduplicate +9. final answer + +### Why this works +Markdown is mostly **topic-organized**, so the agent should first navigate the outline, not the prose. + +--- + +## Workflow B: long JSON memory + +### Example use case +“Summarize all failed tool executions in the last 7 runs and explain recurring causes.” + +### Step-by-step +1. load JSON file +2. inspect top-level keys +3. find `runs`, `events`, `status`, `error`, `timestamp` +4. filter in Python: + - last 7 runs + - failed events only +5. build compact records +6. send batched error records to sub-LM +7. ask sub-LM: + - classify failures + - identify recurring causes + - map to probable remediation +8. aggregate counts in Python +9. final answer + +### Why this works +JSON is better handled as **code-manipulated state**, with the LM used mostly for semantic grouping. + +--- + +## Workflow C: mixed `.md` + `.json` + +### Example use case +“Use architecture docs plus workflow state to explain why the last deployment failed.” + +### Step-by-step +1. search markdown for deployment architecture / rollout / failure handling +2. inspect JSON for recent deployment event logs +3. extract: + - expected behavior from docs + - actual behavior from logs +4. send side-by-side evidence to sub-LM +5. ask for mismatch analysis +6. build final incident explanation + +This is one of the best uses of the paper’s idea: the LM uses external memory as a workspace across heterogeneous sources instead of trying to absorb everything at once. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 8. When to use retrieval vs recursive memory access + +## Use retrieval first when: +- you know likely keywords +- files are numerous +- only a few sections matter + +## Use recursive chunk-by-chunk analysis when: +- answer depends on many sections +- structure matters more than keywords +- you need aggregation over many records +- the task is information-dense + +This mirrors the paper’s distinction between easier sparse tasks and more information-dense tasks where sub-calls and programmatic aggregation matter much more. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 9. A concrete implementation plan + +## Phase 1: basic memory agent +Implement: +- markdown section splitter +- json path accessor +- file manifest +- root loop +- sub-LM batching +- final aggregation buffer + +### Result +A working “memory-aware” agent that does not overflow context. + +--- + +## Phase 2: richer memory tools +Add: +- BM25 or embedding search +- markdown heading search +- JSON schema introspection +- regex over files +- cached chunk summaries +- timestamps and recency weighting + +--- + +## Phase 3: memory compaction +Persist derived memory: +- per-section summary +- per-run summary +- per-error summary +- semantic tags +- embeddings +- entity graph + +Now the agent can reason over: +- raw memory when needed +- compact memory when enough + +--- + +# 10. Sample root-loop pseudocode + +```python +def memory_agent(query, env, root_llm, max_turns=20): + messages = [ + {"role": "system", "content": MEMORY_AGENT_PROMPT}, + {"role": "user", "content": query} + ] + + for _ in range(max_turns): + reply = root_llm(messages) + text = reply["content"] + + if text.startswith("FINAL("): + return text[len("FINAL("):-1] + + code_blocks = extract_repl_blocks(text) + if code_blocks: + outputs = [] + for code in code_blocks: + outputs.append(run_in_memory_env(env, code)) + messages.append({"role": "assistant", "content": text}) + messages.append({"role": "user", "content": "\n\n".join(outputs)}) + else: + messages.append({"role": "assistant", "content": text}) + messages.append({"role": "user", "content": "Continue with memory inspection or return FINAL(...)."}) + + raise RuntimeError("No final answer") +``` + +--- + +# 11. Recommended root prompt for memory gathering + +You can adapt the paper’s RLM prompt into a memory-specific version: + +```text +You are an agent with access to external memory stored in markdown and JSON files. + +Do not try to read everything at once. +First inspect the available files and their structure. +Use code/tools to: +- list files +- inspect markdown outlines +- inspect JSON schemas/keys +- search for likely relevant sections +- read only needed chunks +- batch semantically related chunks into sub-LLM calls +- store intermediate findings in variables + +Prefer programmatic filtering for JSON. +Prefer heading/section-based navigation for Markdown. + +When you have enough evidence, provide FINAL(...) +``` + +That is not a verbatim paper prompt, but it is the most direct practical adaptation of the paper’s mechanism to agent memory. The paper’s own scaffold centers on environment access, chunking, buffering, and recursive sub-calls. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +--- + +# 12. Important pitfalls + +## Pitfall 1: treating memory as one giant prompt +Bad: +- concatenate all docs and logs +- send to LLM + +Good: +- inspect structure first +- read targeted parts only + +## Pitfall 2: using the LM for operations code should do +Bad: +- ask LM to count statuses in JSON + +Good: +- count in Python +- ask LM only to interpret patterns + +## Pitfall 3: overusing sub-calls +The paper notes that too many sub-calls can explode cost and runtime; batching is better. ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +## Pitfall 4: no persistent buffers +If the agent does not store intermediate findings, it will repeatedly rediscover the same information. + +## Pitfall 5: fixed-size chunking for markdown +Heading-aware chunking is much better than arbitrary slices for docs. + +--- + +# 13. Minimal production-ready design + +If you want the simplest robust version, I’d build this: + +## Inputs +- directory of `.md` and `.json` + +## Preprocessing +- markdown heading split +- json schema/path extraction +- embeddings per chunk +- manifest + +## Runtime tools +- list/search/read for markdown +- schema/get/filter for json +- `llm_query()` for semantic interpretation + +## Agent policy +- search → inspect structure → narrow → batch summarize → aggregate + +## Outputs +- final answer +- optional memory summary artifact written back to disk + +--- + +# 14. Best practical use cases + +This pattern is especially strong for: + +- coding agents reading README + config + tool state +- research agents reading notes + results JSON +- ops agents reading incident docs + event logs +- product agents reading specs + task histories +- autonomous workflows that need memory across many prior runs + +--- + +# Final takeaway + +For long `.md` and `.json` files, the paper’s most useful lesson for agent memory is: + +> **Memory should be external, inspectable, structured, and selectively queried—not blindly stuffed into the context window.** ([arxiv.org](https://arxiv.org/abs/2512.24601)) + +So the implementation recipe is: + +1. store raw memory in files +2. derive structural indexes +3. expose memory tools in a persistent environment +4. let the root LM inspect structure first +5. use code for filtering/aggregation +6. use sub-LMs only for semantic interpretation +7. build the final answer from buffers + +If you want, I can now turn this into a **full Python implementation** for: +- **Markdown + JSON memory loader** +- **REPL-style memory environment** +- **sub-LLM querying scaffold** +- **agent loop** \ No newline at end of file diff --git a/setup/gvisor/daemon.json b/environments/gvisor/daemon.json similarity index 100% rename from setup/gvisor/daemon.json rename to environments/gvisor/daemon.json diff --git a/setup/gvisor/daemon.json.new b/environments/gvisor/daemon.json.new similarity index 100% rename from setup/gvisor/daemon.json.new rename to environments/gvisor/daemon.json.new diff --git a/setup/gvisor/install_gvisor.sh b/environments/gvisor/install_gvisor.sh similarity index 100% rename from setup/gvisor/install_gvisor.sh rename to environments/gvisor/install_gvisor.sh diff --git a/setup/images/kalilinux.Dockerfile b/environments/images/kalilinux.Dockerfile similarity index 100% rename from setup/images/kalilinux.Dockerfile rename to environments/images/kalilinux.Dockerfile diff --git a/setup/images/webapp_sec.Dockerfile b/environments/images/webapp_sec.Dockerfile similarity index 100% rename from setup/images/webapp_sec.Dockerfile rename to environments/images/webapp_sec.Dockerfile diff --git a/setup/pgvector/setup_pgvector.sh b/environments/pgvector/setup_pgvector.sh similarity index 100% rename from setup/pgvector/setup_pgvector.sh rename to environments/pgvector/setup_pgvector.sh From 68805bbcdc61bba84820bca3e1d7fa01f60a8db9 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Sat, 28 Mar 2026 21:35:16 +0100 Subject: [PATCH 03/12] Removing sandbox prints resolving a small issue in read tool not converting correctly the lines into int for enumerate. Removing relevance in most agent thoughts return. Those didn't serve any purpose. registry was saving the models api keys as prints. Implementing the tracing to agent thought. --- cli/deadend_cli/types/rpc.ts | 2 - .../agents/components/executor.py | 1 - .../src/deadend_agent/agents/factory.py | 4 +- .../deadend_agent/context/context_engine.py | 19 +-- .../deadend_agent/core_agent/core_agent.py | 149 ++++++++++++++++-- .../deadend_agent/src/deadend_agent/hooks.py | 1 - .../src/deadend_agent/models/registry.py | 1 - .../src/deadend_agent/sandbox/sandbox.py | 30 ++-- .../src/deadend_agent/tools/avfs/read.py | 5 + .../src/deadend_cli/jsonrpc/event_bus.py | 2 - .../src/deadend_cli/jsonrpc/hooks_adapter.py | 2 - .../src/deadend_cli/jsonrpc/rpc_models.py | 3 - 12 files changed, 165 insertions(+), 54 deletions(-) diff --git a/cli/deadend_cli/types/rpc.ts b/cli/deadend_cli/types/rpc.ts index bc11030..1ccde26 100644 --- a/cli/deadend_cli/types/rpc.ts +++ b/cli/deadend_cli/types/rpc.ts @@ -501,12 +501,10 @@ export interface AgentErrorData { * * @property thought - The full reasoning text * @property summary - Condensed version of the thought - * @property relevance - How relevant this thought is to the task (0.0-1.0) */ export interface AgentThoughtData { thought: string; summary?: string; - relevance: number; } /** diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py index 21c92ac..14ebd73 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py @@ -350,7 +350,6 @@ def _add_agent_output_to_context( agent_name=agent_name, thought=thoughts, summary="", # Let context auto-generate summary - relevance=0.9 ) # Log the full agent response - NO TRUNCATION diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py index 7c74630..0d3aa3d 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py @@ -116,11 +116,10 @@ def summarize_agent_thought(raw_thought: str, max_length: int = 200) -> str: class ExtractedThought: """Container for an extracted and summarized agent thought.""" - def __init__(self, agent_name: str, raw_thought: str, summary: str = "", relevance: float = 0.5): + def __init__(self, agent_name: str, raw_thought: str, summary: str = ""): self.agent_name = agent_name self.raw_thought = raw_thought self.summary = summary or summarize_agent_thought(raw_thought) - self.relevance = relevance def to_dict(self) -> dict: """Convert to dictionary for context engine.""" @@ -128,7 +127,6 @@ def to_dict(self) -> dict: "agent_name": self.agent_name, "thought": self.raw_thought, "summary": self.summary, - "relevance": self.relevance } diff --git a/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py b/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py index b54733d..dc39ea5 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/context/context_engine.py @@ -122,13 +122,11 @@ class AgentThought: agent_name: Which agent produced this thought thought: The raw thought/reasoning text summary: A concise summary of the key insight - relevance: How relevant this is for future actions (0.0-1.0) timestamp: When this was recorded """ agent_name: str thought: str summary: str = "" - relevance: float = 0.5 timestamp: float = field(default_factory=time.time) def format_for_context(self) -> str: @@ -334,14 +332,12 @@ def add_thought_simple( agent_name: str, thought: str, summary: str = "", - relevance: float = 0.5 ) -> None: """Convenience method to add an agent thought.""" self.add_thought(AgentThought( agent_name=agent_name, thought=thought, summary=summary, - relevance=relevance )) def was_technique_tested(self, endpoint: str, technique: str) -> bool: @@ -769,12 +765,12 @@ def get_unified_context(self, max_tokens: int = 6000) -> str: lines.append(f"Credential: {af.details.get('username', 'N/A')} / {af.details.get('password', 'N/A')}") sections.append("\n".join(lines)) - # SECTION 7: AGENT INSIGHTS (summarized learnings) + # SECTION 7: AGENT INSIGHTS (recent summarized learnings) if self.thoughts: - high_relevance = [t for t in self.thoughts if t.relevance >= 0.6] - if high_relevance: + recent_thoughts = [t for t in self.thoughts if t.summary or t.thought] + if recent_thoughts: lines = ["## INSIGHTS"] - for t in high_relevance[-5:]: + for t in recent_thoughts[-5:]: summary = t.summary or t.thought[:150] lines.append(f"[{t.agent_name}] {summary}") sections.append("\n".join(lines)) @@ -1493,27 +1489,24 @@ def add_thought( agent_name: str, thought: str, summary: str = "", - relevance: float = 0.5 ) -> None: """Record an agent's reasoning/insight for context. - Thoughts with higher relevance (>= 0.6) are shown to subsequent agents. + Recent thoughts are surfaced to subsequent agents in the INSIGHTS section. Args: agent_name: Which agent produced this thought thought: The raw thought/reasoning text summary: A concise summary of the key insight (auto-generated if empty) - relevance: How relevant this is for future actions (0.0-1.0) Example: context.add_thought( agent_name="shell", thought="The application uses Jinja2 templates based on the error message format.", summary="Application uses Jinja2 templates", - relevance=0.8 ) """ - self.structured.add_thought_simple(agent_name, thought, summary, relevance) + self.structured.add_thought_simple(agent_name, thought, summary) def was_technique_tested(self, endpoint: str, technique: str) -> bool: """Check if a specific technique was already tested on an endpoint. diff --git a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py index 6aa84fd..f428468 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py @@ -51,18 +51,31 @@ # OpenInference span attributes for tool tracing (generic pattern usable by any agent) try: + from openinference.semconv.trace import MessageAttributes as _MsgAttrs from openinference.semconv.trace import SpanAttributes as _SpanAttrs _TOOL_ATTR_KIND = _SpanAttrs.OPENINFERENCE_SPAN_KIND _TOOL_ATTR_NAME = _SpanAttrs.TOOL_NAME _TOOL_ATTR_PARAMS = _SpanAttrs.TOOL_PARAMETERS _TOOL_ATTR_INPUT = _SpanAttrs.INPUT_VALUE _TOOL_ATTR_OUTPUT = _SpanAttrs.OUTPUT_VALUE + _LLM_INPUT_PREFIX = _SpanAttrs.LLM_INPUT_MESSAGES + _LLM_OUTPUT_PREFIX = _SpanAttrs.LLM_OUTPUT_MESSAGES + _LLM_MSG_ROLE = _MsgAttrs.MESSAGE_ROLE + _LLM_MSG_CONTENT = _MsgAttrs.MESSAGE_CONTENT + _LLM_MSG_TOOL_CALLS = _MsgAttrs.MESSAGE_TOOL_CALLS + _LLM_MODEL_NAME = _SpanAttrs.LLM_MODEL_NAME except ImportError: _TOOL_ATTR_KIND = "openinference.span.kind" _TOOL_ATTR_NAME = "tool.name" _TOOL_ATTR_PARAMS = "tool.parameters" _TOOL_ATTR_INPUT = "input.value" _TOOL_ATTR_OUTPUT = "output.value" + _LLM_INPUT_PREFIX = "llm.input_messages" + _LLM_OUTPUT_PREFIX = "llm.output_messages" + _LLM_MSG_ROLE = "message.role" + _LLM_MSG_CONTENT = "message.content" + _LLM_MSG_TOOL_CALLS = "message.tool_calls" + _LLM_MODEL_NAME = "llm.model_name" # Custom attribute for LLM thinking/reasoning content (not part of OpenInference yet) _ATTR_LLM_THINKING = "llm.thinking_content" @@ -389,7 +402,11 @@ async def _run_impl( if thinking_content: assistant_message["thinking_content"] = thinking_content - # Create a child span for this LLM iteration + tool_calls = getattr(choice.message, "tool_calls", None) or [] + llm_trace_full = self._build_llm_trace_output(content, thinking_content, tool_calls) + llm_trace_attr = self._truncate_for_span_attr(llm_trace_full) + + # Create a child span for this LLM iteration (OpenInference LLM + Phoenix .llm_call UI) with self.tracer.start_as_current_span( f"{self.name}.llm_call", attributes={ @@ -397,12 +414,24 @@ async def _run_impl( "llm.iteration": iteration, }, ) as llm_span: - if content: - out_attr = content[:4096] + "..." if len(content) > 4096 else content - llm_span.set_attribute(_TOOL_ATTR_OUTPUT, out_attr) + if self.model: + llm_span.set_attribute(_LLM_MODEL_NAME, self.model) + self._set_llm_input_message_attributes(llm_span, messages) + if llm_trace_attr: + llm_span.set_attribute(_TOOL_ATTR_OUTPUT, llm_trace_attr) + llm_span.set_attribute( + f"{_LLM_OUTPUT_PREFIX}.0.{_LLM_MSG_ROLE}", + "assistant", + ) + llm_span.set_attribute( + f"{_LLM_OUTPUT_PREFIX}.0.{_LLM_MSG_CONTENT}", + llm_trace_attr, + ) + self._set_llm_output_tool_call_attributes(llm_span, tool_calls) if thinking_content: - think_attr = thinking_content[:4096] + "..." if len(thinking_content) > 4096 else thinking_content + think_attr = self._truncate_for_span_attr(thinking_content) llm_span.set_attribute(_ATTR_LLM_THINKING, think_attr) + llm_span.add_event("llm.thinking", {"llm.thinking_content": think_attr}) llm_span.set_status(trace.Status(trace.StatusCode.OK)) # Log thinking content (if any) @@ -442,7 +471,6 @@ async def _run_impl( agent_name=self.name, thought=self._truncate_for_event(thought_text, 3000), summary=self._truncate_for_event(content, 500), - relevance=0.5, ) # Add tool calls if present @@ -526,12 +554,16 @@ async def _run_impl( out_val = output else: out_val = str(output) - out_attr = out_val[:4096] + "..." if len(out_val) > 4096 else out_val - parent_span.set_attribute(_TOOL_ATTR_OUTPUT, out_attr) - # Attach aggregated thinking content to the parent span + parent_trace = self._build_trace_output(out_val, thoughts) + if parent_trace: + parent_span.set_attribute( + _TOOL_ATTR_OUTPUT, + self._truncate_for_span_attr(parent_trace), + ) if thoughts: - think_attr = thoughts[:4096] + "..." if len(thoughts) > 4096 else thoughts + think_attr = self._truncate_for_span_attr(thoughts) parent_span.set_attribute(_ATTR_LLM_THINKING, think_attr) + parent_span.add_event("llm.thinking", {"llm.thinking_content": think_attr}) parent_span.set_status(trace.Status(trace.StatusCode.OK)) return AgentResult( @@ -1312,6 +1344,103 @@ def _extract_thoughts(self, messages: list[dict]) -> str: return result + @staticmethod + def _truncate_for_span_attr(text: str, max_len: int = 4096) -> str: + if len(text) <= max_len: + return text + return text[:max_len] + "..." + + @staticmethod + def _get_attr_or_key(obj: Any, key: str, default: Any = None) -> Any: + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(key, default) + return getattr(obj, key, default) + + @staticmethod + def _message_content_to_text(content: Any) -> str: + if content is None or content == "": + return "" + if isinstance(content, str): + return content + if isinstance(content, list): + parts: list[str] = [] + for block in content: + if isinstance(block, dict): + if block.get("type") == "text" and isinstance(block.get("text"), str): + parts.append(block["text"]) + elif "text" in block: + parts.append(str(block["text"])) + elif isinstance(block, str): + parts.append(block) + return "\n".join(p for p in parts if p) + return str(content) + + def _set_llm_input_message_attributes(self, span: Any, messages: list[dict]) -> None: + """Flatten chat messages into OpenInference ``llm.input_messages.*`` attributes.""" + for index, message in enumerate(messages): + prefix = f"{_LLM_INPUT_PREFIX}.{index}" + role = message.get("role") + if isinstance(role, str) and role: + span.set_attribute(f"{prefix}.{_LLM_MSG_ROLE}", role) + text = self._message_content_to_text(message.get("content", "")) + if text: + span.set_attribute( + f"{prefix}.{_LLM_MSG_CONTENT}", + self._truncate_for_span_attr(text), + ) + + def _set_llm_output_tool_call_attributes(self, span: Any, tool_calls: list[Any]) -> None: + """Flatten tool calls into OpenInference ``llm.output_messages.0.message.tool_calls.*``.""" + base = f"{_LLM_OUTPUT_PREFIX}.0.{_LLM_MSG_TOOL_CALLS}" + for index, tool_call in enumerate(tool_calls): + p = f"{base}.{index}.tool_call" + tid = self._get_attr_or_key(tool_call, "id", "") + if isinstance(tid, str) and tid: + span.set_attribute(f"{p}.id", tid) + function = self._get_attr_or_key(tool_call, "function", None) + fname = self._get_attr_or_key(function, "name", "") + if isinstance(fname, str) and fname: + span.set_attribute(f"{p}.function.name", fname) + fargs = self._get_attr_or_key(function, "arguments", "") + if isinstance(fargs, str) and fargs: + span.set_attribute( + f"{p}.function.arguments", + self._truncate_for_span_attr(fargs), + ) + + @staticmethod + def _build_trace_output(final_text: str, thinking_text: str) -> str: + """Trace-visible text: thinking + final response (matches console-style sections).""" + if final_text and thinking_text: + return f"[Thinking]\n{thinking_text}\n\n[Response]\n{final_text}" + return thinking_text or final_text + + @staticmethod + def _build_llm_trace_output( + content: str, + thinking_content: str, + tool_calls: list[Any], + ) -> str: + """Single-turn LLM trace payload; includes tool-call-only responses.""" + base = CoreAgent._build_trace_output(content, thinking_content) + if base: + return base + if tool_calls: + names: list[str] = [] + for tc in tool_calls: + fn = CoreAgent._get_attr_or_key( + CoreAgent._get_attr_or_key(tc, "function", None), + "name", + "", + ) + if isinstance(fn, str) and fn: + names.append(fn) + if names: + return "[Tool calls] " + ", ".join(names) + return "" + def _get_session_id(self, deps: Any) -> str: """Extract session_id from deps. diff --git a/deadend_cli/deadend_agent/src/deadend_agent/hooks.py b/deadend_cli/deadend_agent/src/deadend_agent/hooks.py index f37337e..94b7e17 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/hooks.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/hooks.py @@ -62,7 +62,6 @@ def emit_agent_thought( agent_name: str, thought: str, summary: Optional[str] = None, - relevance: float = 0.5, ) -> None: """Called when agent reasoning is extracted.""" ... diff --git a/deadend_cli/deadend_agent/src/deadend_agent/models/registry.py b/deadend_cli/deadend_agent/src/deadend_agent/models/registry.py index b43a869..798edb1 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/models/registry.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/models/registry.py @@ -192,7 +192,6 @@ def _initialize_models(self, config: Config): if spec.provider not in self._models: self._models[spec.provider] = [] self._models[spec.provider].append(spec) - logger.info("models init: %s", str(self._models)) # If no providers were found in TOML, registry will simply report no models. diff --git a/deadend_cli/deadend_agent/src/deadend_agent/sandbox/sandbox.py b/deadend_cli/deadend_agent/src/deadend_agent/sandbox/sandbox.py index f0d143b..541153a 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/sandbox/sandbox.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/sandbox/sandbox.py @@ -161,15 +161,13 @@ def start( self.status = SandboxStatus.RUNNING image = self._docker_client.images.get(container_image) - print("Creating container on network: %s", network_name) - # Add host.docker.internal for host access when using non-host networks extra_hosts = None if network_name != "host": extra_hosts = { "host.docker.internal": "host-gateway" # Maps to the host machine } - print("Adding host.docker.internal alias for host access") + # print("Adding host.docker.internal alias for host access") container = self._docker_client.containers.run( image=image, @@ -272,7 +270,7 @@ def execute_command( # Test basic container responsiveness with a simple command health_check = container.exec_run(["/bin/bash", "-c", "echo 'health_check'"]) - print("Health check exit code: %s", health_check.exit_code) + # print("Health check exit code: %s", health_check.exit_code) except Exception as health_exc: print("Container health check failed: %s", health_exc) @@ -287,8 +285,8 @@ def execute_command( try: # Debug: Log the exact command being executed - print("Executing command: %s", ' '.join(shell_command)) - print("Stream mode: %s, Timeout: %s", stream, timeout_seconds) + # print("Executing command: %s", ' '.join(shell_command)) + # print("Stream mode: %s, Timeout: %s", stream, timeout_seconds) if timeout_seconds: # Use threading for timeout implementation @@ -318,9 +316,9 @@ def execute_command( tty=False, # Changed from True to False demux=True ) - print("Command executed, exit code: %s", command_result.exit_code) + # print("Command executed, exit code: %s", command_result.exit_code) (stdout, stderr) = command_result.output - print("Raw stdout length: %d, stderr: %d", len(stdout) if stdout else 0, len(stderr) if stderr else 0) + # print("Raw stdout length: %d, stderr: %d", len(stdout) if stdout else 0, len(stderr) if stderr else 0) result = { "command": command, "exit_code": command_result.exit_code, @@ -330,7 +328,7 @@ def execute_command( "timed_out": False, "execution_time": time.time() - start_time } - print("Decoded stdout length: %d, stderr: %d", len(result['stdout']), len(result['stderr'])) + # print("Decoded stdout length: %d, stderr: %d", len(result['stdout']), len(result['stderr'])) return result @@ -392,7 +390,7 @@ def _execute_with_timeout( def execute_worker(): try: - print("Starting worker thread for command: %s", ' '.join(command if isinstance(command, list) else [command])) + # print("Starting worker thread for command: %s", ' '.join(command if isinstance(command, list) else [command])) start_time = time.time() if stream: command_result = container.exec_run( @@ -402,7 +400,7 @@ def execute_worker(): socket=True, stream=True, ) - print("Streaming command completed") + # print("Streaming command completed") result_container["result"] = { "command": self.last_command, @@ -418,9 +416,9 @@ def execute_worker(): tty=False, # Changed from True to False demux=True ) - print("Command executed, exit code: %s", command_result.exit_code) + # print("Command executed, exit code: %s", command_result.exit_code) (stdout, stderr) = command_result.output - print("Raw stdout length: %d, stderr: %d", len(stdout) if stdout else 0, len(stderr) if stderr else 0) + # print("Raw stdout length: %d, stderr: %d", len(stdout) if stdout else 0, len(stderr) if stderr else 0) result_container["result"] = { "command": self.last_command, "exit_code": command_result.exit_code, @@ -430,12 +428,12 @@ def execute_worker(): "timed_out": False, "execution_time": time.time() - start_time } - print("Decoded stdout length: %d, stderr: %d", len(result_container['result']['stdout']), len(result_container['result']['stderr'])) + # print("Decoded stdout length: %d, stderr: %d", len(result_container['result']['stdout']), len(result_container['result']['stderr'])) except (docker.errors.ContainerError, docker.errors.APIError, OSError) as exc: - print("Docker/system error in worker: %s", exc) + # print("Docker/system error in worker: %s", exc) exception_container["exception"] = exc except Exception as exc: - print("Unexpected error in execute_worker: %s", exc, exc_info=True) + # print("Unexpected error in execute_worker: %s", exc, exc_info=True) exception_container["exception"] = exc thread = threading.Thread(target=execute_worker) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py index a043681..5dfbf9d 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py @@ -26,6 +26,11 @@ async def avfs_read( workspace: str = "workspace", ) -> str: """Read a text file inside the current workspace root with optional 1-based line slicing.""" + start_line = int(start_line) + if end_line is not None: + end_line = int(end_line) + max_chars = int(max_chars) + session_id = _session_id_from_ctx(ctx) target = avfs.resolve(path, session_id=session_id, workspace=workspace) if not target.exists() or not target.is_file(): diff --git a/deadend_cli/src/deadend_cli/jsonrpc/event_bus.py b/deadend_cli/src/deadend_cli/jsonrpc/event_bus.py index a462658..de9445a 100644 --- a/deadend_cli/src/deadend_cli/jsonrpc/event_bus.py +++ b/deadend_cli/src/deadend_cli/jsonrpc/event_bus.py @@ -159,7 +159,6 @@ def emit_agent_thought( agent_name: str, thought: str, summary: Optional[str] = None, - relevance: float = 0.5, ) -> None: """Emit an AGENT_THOUGHT event.""" event = AgentEvent.agent_thought( @@ -167,7 +166,6 @@ def emit_agent_thought( agent_name=agent_name, thought=thought, summary=summary, - relevance=relevance, ) self.publish_sync(event) diff --git a/deadend_cli/src/deadend_cli/jsonrpc/hooks_adapter.py b/deadend_cli/src/deadend_cli/jsonrpc/hooks_adapter.py index 5ad13ff..a3fe8de 100644 --- a/deadend_cli/src/deadend_cli/jsonrpc/hooks_adapter.py +++ b/deadend_cli/src/deadend_cli/jsonrpc/hooks_adapter.py @@ -86,14 +86,12 @@ def emit_agent_thought( agent_name: str, thought: str, summary: Optional[str] = None, - relevance: float = 0.5, ) -> None: self._bus.emit_agent_thought( session_id=session_id, agent_name=agent_name, thought=thought, summary=summary, - relevance=relevance, ) def emit_agent_routed( diff --git a/deadend_cli/src/deadend_cli/jsonrpc/rpc_models.py b/deadend_cli/src/deadend_cli/jsonrpc/rpc_models.py index e2560ff..83e7f39 100644 --- a/deadend_cli/src/deadend_cli/jsonrpc/rpc_models.py +++ b/deadend_cli/src/deadend_cli/jsonrpc/rpc_models.py @@ -163,7 +163,6 @@ class AgentThoughtData(BaseModel): thought: str summary: Optional[str] = None - relevance: float = 0.5 class AgentRoutedData(BaseModel): @@ -419,7 +418,6 @@ def agent_thought( agent_name: str, thought: str, summary: Optional[str] = None, - relevance: float = 0.5, ) -> "AgentEvent": """Create an AGENT_THOUGHT event.""" return cls( @@ -429,7 +427,6 @@ def agent_thought( data=AgentThoughtData( thought=thought, summary=summary, - relevance=relevance, ), ) From affa375271402ba35669bbfdcc1c544fd01f4a28 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Sat, 28 Mar 2026 21:47:39 +0100 Subject: [PATCH 04/12] changing python interpreter sandbox version --- deadend_cli/simple-python-interpreter-sandbox | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deadend_cli/simple-python-interpreter-sandbox b/deadend_cli/simple-python-interpreter-sandbox index 6c99fe5..22cb7a0 160000 --- a/deadend_cli/simple-python-interpreter-sandbox +++ b/deadend_cli/simple-python-interpreter-sandbox @@ -1 +1 @@ -Subproject commit 6c99fe52c39ad40d469c81af112a5cc689b49b9b +Subproject commit 22cb7a0e4d39bdd7cbfb0afaefb30654247363dc From 3c20f64ec2e97bc2da659996b6fc107d5b501aa2 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Sat, 28 Mar 2026 22:18:39 +0100 Subject: [PATCH 05/12] removing pgvector dependency --- .../components/ComponentHealth.tsx | 4 +- cli/deadend_cli/hooks/useComponentHealth.ts | 8 +- cli/deadend_cli/main.tsx | 4 +- cli/deadend_cli/runtime/deadend-rpc-client.ts | 38 ++--- cli/deadend_cli/types/rpc.ts | 2 +- deadend_cli/installer/pyoxidizer.bzl | 1 - deadend_cli/pyproject.toml | 1 - deadend_cli/src/deadend_cli/cli.py | 34 +--- .../src/deadend_cli/component_manager.py | 20 +-- deadend_cli/src/deadend_cli/init.py | 157 +----------------- .../src/deadend_cli/jsonrpc/rpc_server.py | 2 +- deadend_cli/src/deadend_cli/jsonrpc_server.py | 16 +- .../tests/integration/test_rpc_server.py | 12 +- deadend_cli/uv.lock | 14 -- environments/pgvector/setup_pgvector.sh | 14 -- 15 files changed, 55 insertions(+), 272 deletions(-) delete mode 100755 environments/pgvector/setup_pgvector.sh diff --git a/cli/deadend_cli/components/ComponentHealth.tsx b/cli/deadend_cli/components/ComponentHealth.tsx index 405bbe5..e7db27d 100644 --- a/cli/deadend_cli/components/ComponentHealth.tsx +++ b/cli/deadend_cli/components/ComponentHealth.tsx @@ -22,7 +22,7 @@ export interface ComponentHealthProps { // Short display names for components const COMPONENT_NAMES: Record = { docker: "Docker", - pgvector: "pgvector", + rag: "RAG", config: "Config", python_sandbox: "Python", shell_sandbox: "Shell", @@ -31,8 +31,8 @@ const COMPONENT_NAMES: Record = { // Order to display components const COMPONENT_ORDER = [ "docker", - "pgvector", "config", + "rag", "python_sandbox", "shell_sandbox", ]; diff --git a/cli/deadend_cli/hooks/useComponentHealth.ts b/cli/deadend_cli/hooks/useComponentHealth.ts index ed6bbd7..2335a41 100644 --- a/cli/deadend_cli/hooks/useComponentHealth.ts +++ b/cli/deadend_cli/hooks/useComponentHealth.ts @@ -4,7 +4,7 @@ import type { DeadEndRpcClient } from "../runtime/deadend-rpc-client.ts"; export type ComponentName = | "docker" - | "pgvector" + | "rag" | "config" | "python_sandbox" | "shell_sandbox"; @@ -104,8 +104,8 @@ export function useComponentHealth( case "docker": result = await rpcClient.initDocker(); break; - case "pgvector": - result = await rpcClient.initPgvector(); + case "rag": + result = await rpcClient.initRag(); break; case "config": result = await rpcClient.initConfig(); @@ -154,8 +154,8 @@ export function useComponentHealth( const components: ComponentName[] = [ "docker", - "pgvector", "config", + "rag", "python_sandbox", "shell_sandbox", ]; diff --git a/cli/deadend_cli/main.tsx b/cli/deadend_cli/main.tsx index b474ee2..884d3ea 100644 --- a/cli/deadend_cli/main.tsx +++ b/cli/deadend_cli/main.tsx @@ -109,7 +109,7 @@ function App({ cliArgs }: AppProps) { // Initialize all components at once using init_all setInitStatus("Initializing all components..."); - // We can wait longer here for pgvector and the sandbox + // Allow time for sandboxes and Playwright during init_all const initResult = await client.initAll(300000); // Store component results for display @@ -125,7 +125,7 @@ function App({ cliArgs }: AppProps) { } // Check for critical failures (all components required for task execution) - const criticalComponents = ["docker", "config", "model_registry", "pgvector", "shell_sandbox"]; + const criticalComponents = ["docker", "config", "model_registry", "rag", "shell_sandbox"]; const criticalFailures = initResult.failed_components.filter( (c) => criticalComponents.includes(c) ); diff --git a/cli/deadend_cli/runtime/deadend-rpc-client.ts b/cli/deadend_cli/runtime/deadend-rpc-client.ts index 1751495..060c5ee 100644 --- a/cli/deadend_cli/runtime/deadend-rpc-client.ts +++ b/cli/deadend_cli/runtime/deadend-rpc-client.ts @@ -24,7 +24,7 @@ * │ ┌─────────────────────────────────────────────────────────────────────────┐│ * │ │ DeadEndRpcClient ││ * │ │ - runTask() / runTaskWithCallbacks() ││ - * │ │ - healthAll() / initDocker() / initPgvector() / ... ││ + * │ │ - healthAll() / initDocker() / initRag() / ... ││ * │ │ - subscribeEvents() / interrupt() / approve() ││ * │ └────────────────────────────────────────────────────────────────────────┘│ * │ │ │ @@ -67,7 +67,7 @@ * * // Initialize components * await client.initDocker(); - * await client.initPgvector(); + * await client.initRag(); * await client.initShellSandbox(); * * // Run security testing task @@ -301,8 +301,8 @@ export interface DeadEndRpcClientOptions extends StdioRpcClientOptions { * * Before running tasks, components must be initialized in order: * 1. `initDocker()` - Docker daemon connection (required) - * 2. `initPgvector()` - Vector database for RAG (optional) - * 3. `initConfig()` - Load LLM API keys and settings + * 2. `initConfig()` - Load LLM API keys and settings + * 3. `initRag()` - SQLite-backed RAG session manager (after config, for storage paths) * 4. `initShellSandbox()` - Prepare Kali container for shell commands * 5. `initPythonSandbox()` - Start Python interpreter sandbox * 6. `initPlaywright()` - Browser automation (optional) @@ -443,7 +443,7 @@ export class DeadEndRpcClient { * * Returns a comprehensive health report including: * - Docker daemon connectivity - * - pgvector database status + * - RAG (SQLite) session manager status * - Python sandbox process status * - Shell sandbox readiness * - Playwright browser status @@ -466,12 +466,12 @@ export class DeadEndRpcClient { } /** - * Checks pgvector database health. + * Checks RAG (SQLite) session manager health. * - * @returns Promise resolving to HealthResult for pgvector + * @returns Promise resolving to HealthResult for RAG */ - async healthPgvector(): Promise { - const result = await this.client.call("health_pgvector"); + async healthRag(): Promise { + const result = await this.client.call("health_rag"); return result as HealthResult; } @@ -523,17 +523,14 @@ export class DeadEndRpcClient { } /** - * Initializes the pgvector database container. + * Initializes the SQLite-backed RAG session manager. * - * Starts the pgvector container if not running and verifies - * database connectivity. Used for RAG (retrieval-augmented generation). - * - * Requires: initDocker() must be called first + * Prefer calling after `initConfig()` so storage paths from config apply. * * @returns Promise resolving to InitResult with success status */ - async initPgvector(): Promise { - const result = await this.client.call("init_pgvector"); + async initRag(): Promise { + const result = await this.client.call("init_rag"); return result as InitResult; } @@ -611,9 +608,9 @@ export class DeadEndRpcClient { * proper dependency order and provides a comprehensive result. * * Initialization order: - * 1. Docker (required by pgvector and shell_sandbox) - * 2. Config (required by model_registry) - * 3. pgvector (requires Docker) + * 1. Docker (required by shell_sandbox) + * 2. Config (required by model_registry and RAG paths) + * 3. RAG session manager (SQLite, no Docker) * 4. Model Registry (requires Config) * 5. Python sandbox (standalone) * 6. Shell sandbox (requires Docker) @@ -1093,8 +1090,7 @@ export class DeadEndRpcClient { * - Playwright browser * - Python sandbox process * - Shell sandbox containers - * - pgvector database (optional) - * - RAG connector + * - RAG session manager (SQLite) * * @returns Promise resolving to shutdown status for each component */ diff --git a/cli/deadend_cli/types/rpc.ts b/cli/deadend_cli/types/rpc.ts index 1ccde26..10d016a 100644 --- a/cli/deadend_cli/types/rpc.ts +++ b/cli/deadend_cli/types/rpc.ts @@ -645,7 +645,7 @@ export type ComponentStatus = * * Returned by health_* RPC methods to report component status. * - * @property component - Name of the component (docker, pgvector, etc.) + * @property component - Name of the component (docker, rag, config, etc.) * @property healthy - Whether the component is functioning correctly * @property status - Current lifecycle state of the component * @property message - Human-readable status message diff --git a/deadend_cli/installer/pyoxidizer.bzl b/deadend_cli/installer/pyoxidizer.bzl index fc21159..44b303e 100644 --- a/deadend_cli/installer/pyoxidizer.bzl +++ b/deadend_cli/installer/pyoxidizer.bzl @@ -316,7 +316,6 @@ from deadend_cli.jsonrpc_server import main; main() "opentelemetry-api>=1.39.1", "opentelemetry-exporter-otlp>=1.39.1", "opentelemetry-sdk>=1.39.1", - "pgvector>=0.4.1", "playwright>=1.56.0", "prompt-toolkit>=3.0.51", "pydantic>=2.11.5", diff --git a/deadend_cli/pyproject.toml b/deadend_cli/pyproject.toml index 553ec0b..b0026fa 100644 --- a/deadend_cli/pyproject.toml +++ b/deadend_cli/pyproject.toml @@ -26,7 +26,6 @@ dependencies = [ "lxml>=6.0.0", "nest-asyncio>=1.6.0", "numpy>=2.0.0,<2.3.0", - "pgvector>=0.4.1", "playwright>=1.56.0", "prompt-toolkit>=3.0.51", "pydantic>=2.11.5", diff --git a/deadend_cli/src/deadend_cli/cli.py b/deadend_cli/src/deadend_cli/cli.py index be1e021..01a261e 100644 --- a/deadend_cli/src/deadend_cli/cli.py +++ b/deadend_cli/src/deadend_cli/cli.py @@ -14,14 +14,12 @@ import docker import typer -from docker.errors import DockerException from rich.console import Console from deadend_agent import config_setup from deadend_agent.core import start_python_sandbox from .banner import print_banner from .cli_logging import setup_logging -from .init import init_cli_config, check_docker, \ - check_pgvector_container, stop_pgvector_container, setup_pgvector_database +from .init import init_cli_config, check_docker from .chat import Modes, chat_interface from .eval import eval_interface @@ -81,14 +79,6 @@ def chat( ) raise typer.Exit(1) - # Check pgvector database and setup if not running - if not check_pgvector_container(docker_client): - console.print("\n[blue]pgvector database is not running. Setting up...[/blue]") - if not setup_pgvector_database(docker_client): - console.print("\n[red]Failed to setup pgvector database.[/red]") - console.print("Please check Docker logs and try again.") - raise typer.Exit(1) - # Init configuration config = config_setup() log_level_name = str(config.log_level or "INFO").upper() @@ -113,13 +103,6 @@ def chat( finally: if python_process.poll() is None: python_process.terminate() - # Stop pgvector container when chat ends - try: - stop_pgvector_container(docker_client) - except (DockerException, OSError, ConnectionError) as e: - console.print( - f"[yellow]Warning: Could not stop pgvector container: {e}[/yellow]" - ) @app.command() @@ -155,14 +138,6 @@ def eval_agent( ) raise typer.Exit(1) - # Check pgvector database and setup if not running - if not check_pgvector_container(docker_client): - console.print("\n[blue]pgvector database is not running. Setting up...[/blue]") - if not setup_pgvector_database(docker_client): - console.print("\n[red]Failed to setup pgvector database.[/red]") - console.print("Please check Docker logs and try again.") - raise typer.Exit(1) - config = config_setup() log_level_name = str(config.log_level or "INFO").upper() log_level = getattr(logging, log_level_name, logging.INFO) @@ -182,13 +157,6 @@ def eval_agent( finally: if python_process.poll() is None: python_process.terminate() - # Stop pgvector container when chat ends - try: - stop_pgvector_container(docker_client) - except (DockerException, OSError, ConnectionError) as e: - console.print( - f"[yellow]Warning: Could not stop pgvector container: {e}[/yellow]" - ) @app.command() diff --git a/deadend_cli/src/deadend_cli/component_manager.py b/deadend_cli/src/deadend_cli/component_manager.py index 1d7c1fc..a8205a9 100644 --- a/deadend_cli/src/deadend_cli/component_manager.py +++ b/deadend_cli/src/deadend_cli/component_manager.py @@ -48,7 +48,7 @@ def __init__(self): # Component states self.docker_state = ComponentState(name="docker") - self.pgvector_state = ComponentState(name="pgvector") + self.rag_state = ComponentState(name="rag") self.config_state = ComponentState(name="config") self.model_registry_state = ComponentState(name="model_registry") self.python_sandbox_state = ComponentState(name="python_sandbox") @@ -122,18 +122,18 @@ async def init_rag(self) -> InitResult: immediately. """ logger.debug("Initializing RAG session manager...") - self.pgvector_state.status = ComponentStatus.INITIALIZING + self.rag_state.status = ComponentStatus.INITIALIZING try: storage_root = self.config.agents_storage_root if self.config else None self.rag_session_manager = init_rag_session_manager( storage_root=storage_root ) - self.pgvector_state.status = ComponentStatus.READY - self.pgvector_state.metadata["storage_root"] = str( + self.rag_state.status = ComponentStatus.READY + self.rag_state.metadata["storage_root"] = str( self.rag_session_manager._root ) - self.pgvector_state.last_check = datetime.now() + self.rag_state.last_check = datetime.now() logger.debug("RAG session manager initialized at %s", self.rag_session_manager._root) return InitResult( @@ -145,8 +145,8 @@ async def init_rag(self) -> InitResult: ) except Exception as e: logger.error("RAG initialization failed: %s", e) - self.pgvector_state.status = ComponentStatus.ERROR - self.pgvector_state.error_message = str(e) + self.rag_state.status = ComponentStatus.ERROR + self.rag_state.error_message = str(e) return InitResult( success=False, component="rag", @@ -487,7 +487,7 @@ async def health_rag(self) -> HealthResult: ) latency = (time.time() - start_time) * 1000 - self.pgvector_state.last_check = datetime.now() + self.rag_state.last_check = datetime.now() return HealthResult( component="rag", @@ -497,7 +497,7 @@ async def health_rag(self) -> HealthResult: latency_ms=latency, ) except Exception as e: - self.pgvector_state.status = ComponentStatus.UNHEALTHY + self.rag_state.status = ComponentStatus.UNHEALTHY return HealthResult( component="rag", healthy=False, @@ -669,7 +669,7 @@ def get_model(self, provider: str | None = None, model_name: str | None = None): raise RuntimeError( "Model registry not initialized. Call init_model_registry() first." ) - logger.info("models : %s", str(self.model_registry._models)) + if not self.model_registry.has_any_model(): raise RuntimeError( "No LLM model configured. Run `deadend init` to initialize the model configuration." diff --git a/deadend_cli/src/deadend_cli/init.py b/deadend_cli/src/deadend_cli/init.py index bfdb9c4..9583beb 100644 --- a/deadend_cli/src/deadend_cli/init.py +++ b/deadend_cli/src/deadend_cli/init.py @@ -9,13 +9,12 @@ """ import os -import time from pathlib import Path import sys import docker import toml import typer -from docker.errors import DockerException, NotFound +from docker.errors import DockerException from rich.console import Console # Use stderr for console output so that stdout can remain reserved for @@ -46,114 +45,6 @@ def check_docker(client: docker.DockerClient) -> bool: return False -def check_pgvector_container(client: docker.DockerClient) -> bool: - """Check if pgvector container is running. - - Args: - client: Docker client instance - - Returns: - bool: True if pgvector container is running, False otherwise - """ - try: - container = client.containers.get("deadend_pg") - return container.status == "running" - except NotFound: - return False - except DockerException as e: - console.print( - f"[yellow]Warning: Could not check pgvector container status: {e}[/yellow]" - ) - return False - - -def setup_pgvector_database(client: docker.DockerClient) -> bool: - """Setup pgvector database using Docker API. - - Args: - client: Docker client instance - - Returns: - bool: True if setup successful, False otherwise - """ - try: - # Check if container already exists - try: - existing_container = client.containers.get("deadend_pg") - if existing_container.status == "running": - console.print("[green]pgvector database is already running.[/green]") - return True - else: - console.print( - "[yellow]Found existing pgvector container, starting it...[/yellow]" - ) - existing_container.start() - # Wait for container to be ready - time.sleep(5) - console.print("[green]pgvector database started successfully.[/green]") - return True - except NotFound: - pass # Container doesn't exist, create new one - - # Create postgres_data directory in cache if it doesn't exist - cache_dir = Path.home() / ".cache" / "deadend" - postgres_data_dir = cache_dir / "postgres_data" - postgres_data_dir.mkdir(parents=True, exist_ok=True) - - console.print("[blue]Setting up pgvector database...[/blue]") - - # Pull the pgvector image - console.print("Pulling pgvector image...") - client.images.pull("pgvector/pgvector:pg17") - - # Create and run the container - container = client.containers.run( - "pgvector/pgvector:pg17", - name="deadend_pg", - environment={ - "POSTGRES_DB": "codeindexerdb", - "POSTGRES_USER": "postgres", - "POSTGRES_PASSWORD": "postgres", - }, - ports={"5432/tcp": 54320}, - volumes={ - str(postgres_data_dir): { - "bind": "/var/lib/postgresql/data", - "mode": "rw", - } - }, - detach=True, - remove=False, - ) - - # Wait for container to be ready - console.print("Waiting for database to be ready...") - time.sleep(10) - - # Check if container is running - container.reload() - if container.status == "running": - console.print( - "[green]pgvector database setup completed successfully.[/green]" - ) - console.print( - "[blue]Database connection: postgresql://postgres:postgres@localhost:54320/codeindexerdb[/blue]" - ) - return True - else: - console.print( - f"[red]Failed to start pgvector container. Status: {container.status}[/red]" - ) - return False - - except DockerException as e: - console.print(f"[red]Error setting up pgvector database: {e}[/red]") - return False - except (OSError, ConnectionError) as e: - console.print(f"[red]Connection error setting up pgvector: {e}[/red]") - return False - - def pull_sandboxed_kali_image(client: docker.DockerClient) -> bool: """Pull the sandboxed Kali image. @@ -176,36 +67,6 @@ def pull_sandboxed_kali_image(client: docker.DockerClient) -> bool: return False -def stop_pgvector_container(client: docker.DockerClient) -> bool: - """Stop the pgvector container. - - Args: - client: Docker client instance - - Returns: - bool: True if stopped successfully, False otherwise - """ - try: - container = client.containers.get("deadend_pg") - if container.status == "running": - console.print("[blue]Stopping pgvector database...[/blue]") - container.stop() - console.print("[green]pgvector database stopped successfully.[/green]") - return True - else: - console.print("[yellow]pgvector container is not running.[/yellow]") - return True - except NotFound: - console.print("[yellow]pgvector container not found.[/yellow]") - return True - except DockerException as e: - console.print(f"[red]Error stopping pgvector container: {e}[/red]") - return False - except (OSError, ConnectionError) as e: - console.print(f"[red]Connection error stopping pgvector: {e}[/red]") - return False - - def init_cli_config(): """Initialize CLI config by prompting for env vars and saving to cache TOML. @@ -231,16 +92,6 @@ def init_cli_config(): console.print("Please install and start Docker, then run this command again.") raise typer.Exit(1) - # Check and setup pgvector database - if not check_pgvector_container(docker_client): - console.print("\n[blue]pgvector database not found. Setting up...[/blue]") - if not setup_pgvector_database(docker_client): - console.print("\n[red]Failed to setup pgvector database.[/red]") - console.print("Please check Docker logs and try again.") - raise typer.Exit(1) - else: - console.print("[green]pgvector database is already running.[/green]") - # Pull sandboxed Kali image console.print("\n[blue]Setting up sandboxed Kali image...[/blue]") if not pull_sandboxed_kali_image(docker_client): @@ -305,9 +156,7 @@ def init_cli_config(): "LOCAL_MODEL": os.getenv("LOCAL_MODEL", "Kimi-K2-Thinking"), "LOCAL_BASE_URL": os.getenv("LOCAL_BASE_URL", ""), "EMBEDDING_MODEL": os.getenv("EMBEDDING_MODEL", ""), - "DB_URL": os.getenv( - "DB_URL", "postgresql://postgres:postgres@localhost:54320/codeindexerdb" - ), + "DB_URL": os.getenv("DB_URL", ""), "ZAP_PROXY_API_KEY": os.getenv("ZAP_PROXY_API_KEY", ""), "APP_ENV": os.getenv("APP_ENV", "development"), "LOG_LEVEL": os.getenv("LOG_LEVEL", "INFO"), @@ -329,7 +178,7 @@ def init_cli_config(): ("LOCAL_MODEL", False, "Local model name (e.g., Kimi-K2-Thinking)"), ("LOCAL_BASE_URL", False, "Local model base URL (e.g., http://localhost:8000/v1)"), ("EMBEDDING_MODEL", False, "Embedding model (optional, for RAG features)"), - ("DB_URL", False, "Database URL (optional)"), + ("DB_URL", False, "Legacy PostgreSQL URL (optional; RAG uses SQLite per session)"), ("ZAP_PROXY_API_KEY", True, "ZAP Proxy API key (optional, for security testing)"), ("APP_ENV", False, "Application environment"), ("LOG_LEVEL", False, "Log level (INFO, DEBUG, etc.)"), diff --git a/deadend_cli/src/deadend_cli/jsonrpc/rpc_server.py b/deadend_cli/src/deadend_cli/jsonrpc/rpc_server.py index d00068f..954c303 100644 --- a/deadend_cli/src/deadend_cli/jsonrpc/rpc_server.py +++ b/deadend_cli/src/deadend_cli/jsonrpc/rpc_server.py @@ -5,7 +5,7 @@ """JSON-RPC server over stdio for communicating with other front-end components. This server supports: -- Component initialization (Docker, pgvector, config, sandboxes, Playwright) +- Component initialization (Docker, RAG/SQLite, config, sandboxes, Playwright) - Health checks for all components - Event streaming for agent/tool execution - Approval workflow for dangerous tool calls diff --git a/deadend_cli/src/deadend_cli/jsonrpc_server.py b/deadend_cli/src/deadend_cli/jsonrpc_server.py index 596715b..8560f37 100644 --- a/deadend_cli/src/deadend_cli/jsonrpc_server.py +++ b/deadend_cli/src/deadend_cli/jsonrpc_server.py @@ -168,14 +168,14 @@ async def init_docker( result = await component_manager.init_docker() return result.model_dump() - @server.add_method("init_pgvector") - async def init_pgvector( + @server.add_method("init_rag") + async def init_rag( _request_id: Any, _params: Dict[str, Any], component_manager: ComponentManager ) -> Dict[str, Any]: - """Initialize pgvector database.""" - result = await component_manager.init_pgvector() + """Initialize SQLite-backed RAG session manager.""" + result = await component_manager.init_rag() return result.model_dump() @server.add_method("init_config") @@ -239,14 +239,14 @@ async def health_docker( result = await component_manager.health_docker() return result.model_dump() - @server.add_method("health_pgvector") - async def health_pgvector( + @server.add_method("health_rag") + async def health_rag( _request_id: Any, _params: Dict[str, Any], component_manager: ComponentManager ) -> Dict[str, Any]: - """Check pgvector health.""" - result = await component_manager.health_pgvector() + """Check RAG (SQLite) session manager health.""" + result = await component_manager.health_rag() return result.model_dump() @server.add_method("health_python_sandbox") diff --git a/deadend_cli/tests/integration/test_rpc_server.py b/deadend_cli/tests/integration/test_rpc_server.py index 21bd653..6e465fe 100644 --- a/deadend_cli/tests/integration/test_rpc_server.py +++ b/deadend_cli/tests/integration/test_rpc_server.py @@ -427,9 +427,9 @@ async def test_health_all(self, rpc_client): @pytest.mark.asyncio @pytest.mark.slow - async def test_health_pgvector(self, rpc_client): - """Test health_pgvector method.""" - await rpc_client.send_request("health_pgvector") + async def test_health_rag(self, rpc_client): + """Test health_rag method.""" + await rpc_client.send_request("health_rag") response = await rpc_client.read_response(timeout=10.0) assert response["jsonrpc"] == "2.0" @@ -702,9 +702,9 @@ async def test_init_docker(self, rpc_client): @pytest.mark.asyncio @pytest.mark.slow @pytest.mark.docker - async def test_init_pgvector(self, rpc_client): - """Test init_pgvector method (requires Docker).""" - await rpc_client.send_request("init_pgvector") + async def test_init_rag(self, rpc_client): + """Test init_rag method.""" + await rpc_client.send_request("init_rag") response = await rpc_client.read_response(timeout=60.0) assert response["jsonrpc"] == "2.0" diff --git a/deadend_cli/uv.lock b/deadend_cli/uv.lock index 928be00..677244d 100644 --- a/deadend_cli/uv.lock +++ b/deadend_cli/uv.lock @@ -1110,7 +1110,6 @@ dependencies = [ { name = "openinference-instrumentation" }, { name = "openinference-instrumentation-instructor" }, { name = "openinference-instrumentation-litellm" }, - { name = "pgvector" }, { name = "playwright" }, { name = "prompt-toolkit" }, { name = "pydantic" }, @@ -1180,7 +1179,6 @@ requires-dist = [ { name = "openinference-instrumentation", specifier = ">=0.1.44" }, { name = "openinference-instrumentation-instructor", specifier = ">=0.1.13" }, { name = "openinference-instrumentation-litellm", specifier = ">=0.1.29" }, - { name = "pgvector", specifier = ">=0.4.1" }, { name = "playwright", specifier = ">=1.56.0" }, { name = "prompt-toolkit", specifier = ">=3.0.51" }, { name = "pydantic", specifier = ">=2.11.5" }, @@ -3336,18 +3334,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9a/70/875f4a23bfc4731703a5835487d0d2fb999031bd415e7d17c0ae615c18b7/pathvalidate-3.3.1-py3-none-any.whl", hash = "sha256:5263baab691f8e1af96092fa5137ee17df5bdfbd6cff1fcac4d6ef4bc2e1735f", size = 24305, upload-time = "2025-06-15T09:07:19.117Z" }, ] -[[package]] -name = "pgvector" -version = "0.4.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/44/43/9a0fb552ab4fd980680c2037962e331820f67585df740bedc4a2b50faf20/pgvector-0.4.1.tar.gz", hash = "sha256:83d3a1c044ff0c2f1e95d13dfb625beb0b65506cfec0941bfe81fd0ad44f4003", size = 30646, upload-time = "2025-04-26T18:56:37.151Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bf/21/b5735d5982892c878ff3d01bb06e018c43fc204428361ee9fc25a1b2125c/pgvector-0.4.1-py3-none-any.whl", hash = "sha256:34bb4e99e1b13d08a2fe82dda9f860f15ddcd0166fbb25bffe15821cbfeb7362", size = 27086, upload-time = "2025-04-26T18:56:35.956Z" }, -] - [[package]] name = "platformdirs" version = "4.5.0" diff --git a/environments/pgvector/setup_pgvector.sh b/environments/pgvector/setup_pgvector.sh deleted file mode 100755 index 744d73d..0000000 --- a/environments/pgvector/setup_pgvector.sh +++ /dev/null @@ -1,14 +0,0 @@ -#!/bin/bash - -# PG vector setup for RAG -mkdir postgres_data -docker run --rm \ - -e POSTGRES_DB=codeindexerdb \ - -e POSTGRES_USER=postgres \ - -e POSTGRES_PASSWORD=postgres \ - -d \ - --name deadend_pg \ - -p 54320:5432 \ - -v ./postgres_data:/var/lib/postgresql/data \ - pgvector/pgvector:pg17 - From 75a8bf7b7a66d76a875717b387da216f7e629b1b Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Sat, 28 Mar 2026 22:40:50 +0100 Subject: [PATCH 06/12] clean up --- .../src/deadend_agent/agents/reporter.py | 12 +++++++++++- .../src/deadend_agent/agents/validator.py | 6 +++--- deadend_cli/src/deadend_cli/eval.py | 3 +-- 3 files changed, 15 insertions(+), 6 deletions(-) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py index 7bdd197..567402a 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py @@ -37,7 +37,17 @@ def __init__(self, model, deps_type, tools, validation_type: str | None, validat output_type=ReporterOutput, tools=[] ) - async def run(self, prompt, deps, message_history, usage, usage_limits, deferred_tool_results=None): + async def run( + self, + prompt, + deps, + message_history, + usage, + usage_limits, + deferred_tool_results=None, + *args, + **kwargs + ): return await super().run( prompt=prompt, deps=deps, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/validator.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/validator.py index 4292118..1c21db9 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/validator.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/validator.py @@ -10,7 +10,6 @@ """ from typing import Dict from pydantic import BaseModel -from pydantic_ai import AgentRunResult from deadend_prompts import render_agent_instructions from .factory import AgentRunner @@ -41,7 +40,6 @@ def __init__(self, model, deps_type, tools, available_agents: Dict[str, str]): tools=[] ) - async def run( self, prompt, @@ -49,7 +47,9 @@ async def run( message_history, usage, usage_limits, - deferred_tool_results + deferred_tool_results=None, + *args, + **kwargs ): return await super().run( prompt=prompt, diff --git a/deadend_cli/src/deadend_cli/eval.py b/deadend_cli/src/deadend_cli/eval.py index a5c22c1..4bbcc30 100644 --- a/deadend_cli/src/deadend_cli/eval.py +++ b/deadend_cli/src/deadend_cli/eval.py @@ -95,7 +95,6 @@ async def eval_interface( await eval_deadend_agent( model=model_registry.get_model(provider=providers[0]), embedder_client=embedder_client, - # evaluators=[CtfEvaluator], code_indexer_db=rag_db, sandbox=sandbox, eval_metadata=eval_metadata, @@ -107,4 +106,4 @@ async def eval_interface( output_report="./", hard_prompt=False ) - # for model in models: + From 96943bdd3d744eda9a4427fbe10124eaaf88c97f Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Sun, 29 Mar 2026 23:03:29 +0200 Subject: [PATCH 07/12] Integrating validation strategy. Some issues around the write on the avfs write tooling got stuck. It needed changes to a wrapper tooling to not let the model decide the workspace, which was a pretty idiotic thing to do. I mostly blame claude, but I need to move fast, so fuck this shit. Added changes to the instructions tools descriptions according to the tools added. the validation for now will be kept as is until after the platform launch. --- deadend_cli/README.md | 112 ++++ .../src/deadend_agent/agents/__init__.py | 21 +- .../src/deadend_agent/agents/architecture.py | 495 ++++++++---------- .../agents/components/executor.py | 40 +- .../components/validation_strategies.py | 455 ++++++++++++++++ .../agents/components/validator.py | 123 ----- .../agents/generic_agents/memory_agent.py | 23 +- .../src/deadend_agent/agents/reporter.py | 211 ++++++-- .../config/validation.default.yaml | 62 +++ .../deadend_agent/core_agent/core_agent.py | 20 +- .../src/deadend_agent/deadend_agent.py | 81 +-- .../src/deadend_agent/tools/__init__.py | 38 +- .../src/deadend_agent/tools/avfs/__init__.py | 40 +- .../src/deadend_agent/tools/avfs/list.py | 155 +++++- .../src/deadend_agent/tools/avfs/read.py | 107 +++- .../src/deadend_agent/tools/avfs/write.py | 42 +- .../deadend_agent/tests/rlm/test_avfs.py | 34 +- .../rlm/test_deadend_agent_avfs_startup.py | 16 +- .../tests/rlm/test_memory_avfs.py | 39 +- .../deadend_eval/src/deadend_eval/eval.py | 37 +- .../_shared/_anti_fabrication.jinja2 | 33 ++ .../_shared/_memory_summary.jinja2 | 4 +- .../exploit_web.instructions.jinja2 | 19 +- .../deadend_prompts/judge.instructions.jinja2 | 33 +- .../memory.instructions.jinja2 | 10 +- .../planner.instructions.jinja2 | 15 +- .../python_interpreter.instructions.jinja2 | 90 +--- .../recon_threatmodel.instructions.jinja2 | 14 +- .../reporter.instructions.jinja2 | 17 +- .../requester.instructions.jinja2 | 90 +--- .../deadend_prompts/shell.instructions.jinja2 | 19 +- .../supervisor.instructions.jinja2 | 9 +- .../grep_memory_files.description.jinja2 | 1 + .../grep_workspace_files.description.jinja2 | 1 + .../list_memory_files.description.jinja2 | 1 + .../list_workspace_files.description.jinja2 | 1 + .../tools/read_memory_file.description.jinja2 | 1 + .../read_workspace_file.description.jinja2 | 1 + .../write_memory_file.description.jinja2 | 1 + .../write_workspace_file.description.jinja2 | 1 + .../webapp_analyzer.instructions.jinja2 | 19 +- .../webapp_recon.instructions.jinja2 | 19 +- deadend_cli/src/deadend_cli/chat.py | 11 +- 43 files changed, 1656 insertions(+), 905 deletions(-) create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py delete mode 100644 deadend_cli/deadend_agent/src/deadend_agent/agents/components/validator.py create mode 100644 deadend_cli/deadend_agent/src/deadend_agent/config/validation.default.yaml create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_anti_fabrication.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_memory_files.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_workspace_files.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_memory_files.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_workspace_files.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_memory_file.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_workspace_file.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_memory_file.description.jinja2 create mode 100644 deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_workspace_file.description.jinja2 diff --git a/deadend_cli/README.md b/deadend_cli/README.md index 96973bf..d1dc470 100644 --- a/deadend_cli/README.md +++ b/deadend_cli/README.md @@ -159,6 +159,118 @@ The agent uses a two-phase approach (reconnaissance → exploitation) with a sup --- +## Validation Configuration + +The agent uses a composable validation system to determine when the root goal of an assessment has been achieved. Configuration is driven by a YAML file at `~/.cache/deadend/validation.yaml`. + +### How It Works + +After every supervisor execution, a **validation gate** runs a chain of strategies in order. The first strategy that returns `stop: true` triggers a report and exits the loop. If no strategy stops, the ADaPT policy (expand/refine/fail) continues as normal. + +**Available strategies:** + +| Strategy | Cost | What it does | +|----------|------|-------------| +| `flag` | Zero (regex) | Scans proofs, summaries, and context for a token matching a configurable regex pattern | +| `judge` | 1 LLM call | Agent that evaluates the full execution trace against the root goal. Self-throttles when no new evidence has appeared | + +### Configuration File + +Create `~/.cache/deadend/validation.yaml`. A reference file with all options is at `deadend_agent/src/deadend_agent/config/validation.default.yaml`. + +### Examples + +**CTF with `FLAG{}` tokens** (default if no file exists): + +```yaml +validation_format: "FLAG{}" +validation_type: "flag" +strategies: + - name: flag + pattern: "FLAG\\{[^}]+\\}" + - name: judge +``` + +The `flag` strategy runs first (free regex check). If no match, the `judge` LLM evaluates whether the goal is done. + +**HackTheBox:** + +```yaml +validation_format: "HTB{}" +validation_type: "flag" +strategies: + - name: flag + pattern: "HTB\\{[^}]+\\}" + - name: judge + validation_format: "HTB{}" +``` + +**picoCTF:** + +```yaml +validation_format: "picoCTF{}" +validation_type: "flag" +strategies: + - name: flag + pattern: "picoCTF\\{[^}]+\\}" + - name: judge + validation_format: "picoCTF{}" +``` + +**Recon / security assessment** (no flag to find): + +```yaml +validation_type: "security assessment" +strategies: + - name: judge +``` + +No `flag` strategy — the LLM judge evaluates whether the recon goal (e.g., "map the attack surface") is satisfied based on accumulated evidence. + +**Flag-only (fastest, no LLM judge):** + +```yaml +validation_format: "FLAG{}" +strategies: + - name: flag +``` + +Only regex matching, no LLM call at all. Cheapest option for CTFs where the flag format is known. + +### Configuration Reference + +**Top-level fields:** + +| Field | Type | Description | +|-------|------|-------------| +| `validation_format` | `string \| null` | Token format shown in agent prompts (e.g., `"FLAG{}"`, `"HTB{}"`). Set to `null` for assessments without tokens | +| `validation_type` | `string \| null` | Type label for the judge prompt (e.g., `"flag"`, `"security assessment"`) | +| `strategies` | `list` | Ordered list of strategy configurations | + +**Per-strategy fields:** + +| Field | Strategy | Description | +|-------|----------|-------------| +| `name` | all | Strategy name: `"flag"` or `"judge"` | +| `pattern` | `flag` | Regex pattern (default: `FLAG\{[^}]+\}`) | +| `validation_type` | `judge` | Override top-level `validation_type` for this strategy | +| `validation_format` | `judge` | Override top-level `validation_format` for this strategy | + +### Programmatic Override + +Pass a custom config path when constructing the agent: + +```python +agent = DeadEndAgent( + session_id=session_id, + model=model, + available_agents=agents, + validation_config_path="/path/to/custom/validation.yaml", +) +``` + +--- + ## Benchmark Results Evaluated on XBOW's 104-challenge validation suite (black-box mode, January 2026): diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py index efecc1f..7f5419e 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/__init__.py @@ -4,7 +4,6 @@ from .planner import Planner, PlannerAgent, PlannerOutput, RagDeps from .supervisor_agent import SupervisorAgent, SupervisorOutput -from .judge import JudgeAgent, JudgeOutput from .factory import AgentRunner, AgentOutput from .generic_agents.shell_agent import ShellAgent, ShellOutput from .generic_agents.python_interpreter_agent import PythonInterpreterAgent, PythonInterpreterOutput @@ -13,15 +12,13 @@ from .generic_agents.memory_agent import MemoryAgent __all__ = [ - "AgentRunner", "AgentOutput", - "Planner", "PlannerAgent", "PlannerOutput", "RagDeps", - "SupervisorAgent", "SupervisorOutput", - "JudgeOutput", "JudgeAgent", - # Generic agents - "ShellAgent", "ShellOutput", - "PythonInterpreterAgent", "PythonInterpreterOutput", - "RequesterAgent", "RequesterOutput", - "WebAppAnalyzerAgent", - "MemoryAgent" - + "AgentRunner", "AgentOutput", + "Planner", "PlannerAgent", "PlannerOutput", "RagDeps", + "SupervisorAgent", "SupervisorOutput", + # Generic agents + "ShellAgent", "ShellOutput", + "PythonInterpreterAgent", "PythonInterpreterOutput", + "RequesterAgent", "RequesterOutput", + "WebAppAnalyzerAgent", + "MemoryAgent", ] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py index c7e46b8..6398c07 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py @@ -1,20 +1,17 @@ +from __future__ import annotations -import re -from typing import Any, Literal, AsyncGenerator, Tuple +from typing import Any, Literal, AsyncGenerator from uuid import UUID -from uuid import UUID, uuid4 -from deadend_agent.agents.exploit_web_agent import ExploitInfo, ExploitOutput -from deadend_agent.agents.recon_threatmodel_agent import GeneralInfoOutput, ThreatModelOutput -from pydantic import BaseModel, Field -from pydantic_ai import DeferredToolResults, RunContext -from pydantic_ai.exceptions import UsageLimitExceeded + from pydantic_ai.usage import RunUsage, UsageLimits -from deadend_agent.agents.components.planner import TaskNode -from deadend_agent.utils.structures import TaskPlanner + from deadend_agent.agents.components.executor import AgentExecutor, ResultEvent, LogEvent -from deadend_agent.agents.components.planner import Planner -from deadend_agent.agents.components.validator import Validator +from deadend_agent.agents.components.planner import Planner, TaskNode +from deadend_agent.agents.components.validation_strategies import ValidationGate, ValidationVerdict +from deadend_agent.agents.reporter import ReporterAgent from deadend_agent.context import ContextEngine +from deadend_agent.logging import logger +from deadend_agent.utils.structures import TaskPlanner def _format_dict_for_context(data: dict[str, Any]) -> str: @@ -63,110 +60,113 @@ class ADaPTAgent: def __init__( self, session_id: UUID, + agent_id: UUID, context: ContextEngine, executor: AgentExecutor, planner: Planner, - validator: Validator, - max_depth: int = 3 + validation_gate: ValidationGate, + reporter: ReporterAgent, + max_depth: int = 3, ): """Initialize the ADaPT agent. Args: - session_id: Unique identifier for this ADaPT session - executor: AgentExecutor instance for executing tasks - planner: Planner instance for decomposing tasks - validator: Validator instance for validating task completion - max_depth: Maximum depth for task decomposition (default: 3) + session_id: Unique identifier for this ADaPT session. + agent_id: Local agent ID used for AVFS workspace mounts. + context: Shared context engine across all agents. + executor: AgentExecutor instance for executing tasks. + planner: Planner instance for decomposing tasks. + validation_gate: Composable gate that checks whether the + root goal is satisfied after every supervisor return. + reporter: ReporterAgent that writes an MD report on stop. + max_depth: Maximum depth for task decomposition (default: 3). """ self.session = session_id + self.agent_id = agent_id self.max_depth = max_depth self.executor = executor self.planner = planner - self.validator = validator + self.validation_gate = validation_gate + self.reporter = reporter self.context = context - # Track attempted tasks to prevent redundant retries - # Key: task hash, Value: dict with attempt count and best confidence + # Track attempted tasks to prevent redundant retries. + # Key: task hash, Value: dict with attempt count and best confidence. self._attempted_tasks: dict[int, dict[str, Any]] = {} async def _solve( self, node: TaskNode, depth: int, - exit_strategy: str, - exit_loop: bool + exit_loop: bool, ) -> AsyncGenerator[str | dict[str, Any], None]: """Recursively solve a task node using the ADaPT algorithm. - This is the core recursive method that implements the ADaPT algorithm: - 1. Execute the task and get confidence score - 2. Apply policy to determine next action (fail, expand, refine, validate) - 3. If expanding, decompose into subtasks and recursively solve each - 4. If validating, verify task completion + After every supervisor return the validation gate is consulted. + If the gate says stop (root goal achieved), write a report and + propagate exit_loop. Otherwise fall through to the ADaPT policy + (fail / expand / refine). Args: - node: The TaskNode to solve - depth: Current depth in the decomposition tree - - Note: - Tasks that exceed max_depth will be marked as "aborted:max_depth". - The method modifies the node's status and confidence_score in place. - This is implemented as an async generator to stream intermediate log entries - upstream while still performing recursive execution. + node: The TaskNode to solve. + depth: Current depth in the decomposition tree. + exit_loop: Whether an ancestor already requested termination. """ def emit(message: str) -> str: - """Surface message to callers without adding to context (avoid duplication).""" return message - # Check max depth + # --- Guard: max depth --- if depth > self.max_depth: node.status = "aborted:max_depth" node.confidence_score = 0.5 yield emit(f"[ADAPT] Aborted task '{node.task}' at depth {depth} (max_depth={self.max_depth})") return - # Track attempts for this task to prevent infinite loops + # --- Guard: duplicate / exhausted tasks --- task_hash = hash(node.task) if task_hash not in self._attempted_tasks: self._attempted_tasks[task_hash] = {"attempts": 0, "best_confidence": 0.0} task_record = self._attempted_tasks[task_hash] - # Check if task was already completed with high confidence if task_record["best_confidence"] >= self.VALIDATE_THRESHOLD: node.status = "completed" node.confidence_score = task_record["best_confidence"] yield emit(f"[SKIP] Task already completed with {task_record['best_confidence']:.2f} confidence") return - # Check max attempts if task_record["attempts"] >= self.MAX_TASK_ATTEMPTS: node.status = "failed:max_attempts" node.confidence_score = task_record["best_confidence"] yield emit(f"[FAIL] Task exceeded max attempts ({self.MAX_TASK_ATTEMPTS}): '{node.task[:50]}...'") return - # Increment attempt counter task_record["attempts"] += 1 - - # Clear current task log for fresh context self.context.clear_current_task_log() - # Use a local variable to track exit_loop since the parameter can't be modified should_exit = exit_loop - while not should_exit or node.status != 'completed': - # We need to add that this loop should run while we still have the task running - # and not comple otherwise we should exit_loop too. - # Build UNIFIED context for executor (same as all other agents) - # This ensures router, executor, validator all see the same information + # Accumulate a log of supervisor ↔ subagent interactions across + # iterations so the next supervisor call knows what was already tried. + interaction_history: list[str] = [] + + while not should_exit or node.status != "completed": + # ----- 1. Execute supervisor ----- tasks_context = self.context.get_tasks(depth=0, include_goal=False) unified_context = self.context.get_unified_context(max_tokens=6000) - agent_context = f"{unified_context}\n\n{tasks_context}" - # First here, we should give the executor the right task. which means the right - # task with the right context to achieve this task.the supervisor is actually - # more of a subagent. - executor_stream = self.executor.execute_supervisor(task_node=node, agent_context=agent_context) + + if interaction_history: + history_block = ( + "## Previous Supervisor Iterations (DO NOT repeat these actions)\n" + + "\n".join(interaction_history) + ) + agent_context = f"{unified_context}\n\n{history_block}\n\n{tasks_context}" + else: + agent_context = f"{unified_context}\n\n{tasks_context}" + + executor_stream = self.executor.execute_supervisor( + task_node=node, agent_context=agent_context, + ) confidence_score: float | None = None task_achieved: bool = False @@ -179,36 +179,31 @@ def emit(message: str) -> str: confidence_score = event.confidence_score new_context = event.context - # Extract SupervisorOutput fields from context task_achieved = new_context.get("task_achieved", False) detailed_summary = new_context.get("detailed_summary", "") proofs = new_context.get("proofs", "") - # Update node status based on task_achieved if task_achieved: node.status = "completed" node.confidence_score = confidence_score self.context.mark_task_completed(node.task, confidence_score) yield emit(f"[SUPERVISOR] Task achieved: {node.task[:50]}...") - # Add detailed summary to context if detailed_summary: self.context.add_agent_response( f"[Supervisor] {detailed_summary}", - skip_structured=False + skip_structured=False, ) - # Add proofs to context as discovered facts if proofs: self.context.add_discovered_fact( category="proof", key=f"proof_{node.task[:30]}", value=proofs, confidence=confidence_score, - actionable=not task_achieved + actionable=not task_achieved, ) - # Log task status status_str = "ACHIEVED" if task_achieved else "IN PROGRESS" self.context.structured.append_to_log( f"[Supervisor] Task: {status_str} | Confidence: {confidence_score:.2f}" @@ -216,7 +211,6 @@ def emit(message: str) -> str: break elif isinstance(event, LogEvent): - # Only add to structured log, not full workflow_context (reduces duplication) self.context.structured.append_to_log(event.message) yield emit(event.message) else: @@ -225,103 +219,133 @@ def emit(message: str) -> str: if confidence_score is None or new_context is None: raise RuntimeError("AgentExecutor did not produce a result.") - # Update best confidence for this task if confidence_score > task_record["best_confidence"]: task_record["best_confidence"] = confidence_score - # Emit summary for streaming output + # Record this iteration so the next supervisor call knows what + # was already attempted and does not repeat the same actions. + iteration_entry = ( + f"--- Iteration {len(interaction_history) + 1} ---\n" + f"Task: {node.task}\n" + f"Result: {'achieved' if task_achieved else 'not achieved'} " + f"| confidence={confidence_score:.2f}\n" + f"Summary: {detailed_summary}\n" + ) + if proofs: + iteration_entry += f"Evidence: {proofs}\n" + # Include the subagent interaction log captured by emit() inside + # execute_supervisor tool calls (stored in context["log"]). + exec_log = new_context.get("log", "") + if exec_log: + iteration_entry += f"Subagent calls:\n{exec_log}\n" + interaction_history.append(iteration_entry) + if detailed_summary: yield emit(f"[RESULT] {detailed_summary[:200]}") - # If supervisor confirmed task achieved, check for flag and move on + # ----- 2. Validation gate ----- + # Only run the full gate (which may include expensive LLM judge) + # when there is reason to believe the root goal might be done: + # either the supervisor says the subtask is achieved, or + # confidence is high enough to warrant checking. + supervisor_output = { + "task_achieved": task_achieved, + "detailed_summary": detailed_summary, + "proofs": proofs, + "confidence_score": confidence_score, + } + + should_validate = ( + task_achieved + or confidence_score >= self.VALIDATE_THRESHOLD + ) + + if should_validate: + validation_context = self.context.get_unified_context(max_tokens=8000) + verdict = await self.validation_gate.check( + output=supervisor_output, + root_goal=self.context.final_goal, + context=validation_context, + ) + + if verdict.stop: + node.status = "completed" + node.confidence_score = verdict.confidence + self.context.mark_task_completed(node.task, verdict.confidence) + yield emit(f"[VALIDATION] Root goal achieved (confidence={verdict.confidence:.2f})") + + # Write report via the reporter agent. + await self._write_report(verdict) + + if verdict.token: + yield {"validation_token": verdict.token} + yield {"exit_loop": True} + return + + # ----- 3. If subtask done but root goal not yet, continue ----- if task_achieved: - yield emit(f"[SUPERVISOR] Task completed: {node.task[:50]}...") - # Check if proofs contain a flag - if so, exit the loop - if proofs and "FLAG{" in proofs.upper(): - flag_match = re.search(r'FLAG\{[^}]+\}', proofs, re.IGNORECASE) - if flag_match: - yield {'validation_token': flag_match.group(0)} - yield {'exit_loop': True} - return # Task achieved, move to next task + yield emit(f"[SUPERVISOR] Subtask completed: {node.task[:50]}...") + return + # ----- 4. ADaPT policy (expand / refine / fail) ----- decision = self._policy(confidence_score) - try: - print(f"task: {node.task[:50]}... decision: {decision}, \ - confidence: {confidence_score:.2f}") - except (BlockingIOError, OSError): - pass + logger.debug( + "task: %s decision: %s confidence: %.2f", + node.task[:50], decision, confidence_score, + ) - # If decision fails <20% if decision == "fail": node.status = "failed" self.context.update_task_status(node.task, "failed", confidence_score) - yield emit(f"[POLICY] Task '{node.task[:50]}...' \ - failed with confidence {confidence_score:.2f}") + yield emit(f"[POLICY] Task '{node.task[:50]}...' failed with confidence {confidence_score:.2f}") return - # If the decision is validate >80% elif decision == "validate": - node.status, validation_token = await self._validate(node) - yield emit(f"[POLICY] Validation completed for \ - '{node.task[:50]}...' with status {node.status}") - # Update task status in context - if node.status == "completed": - self.context.mark_task_completed(node.task, node.confidence_score) - else: - self.context.update_task_status(node.task, node.status, node.confidence_score) - if len(validation_token) > 1: - yield {'validation_token': validation_token} - yield {'exit_loop': True} + # High subtask confidence — mark subtask done, loop will + # re-check root goal on next iteration if there's more work. + node.status = "completed" + node.confidence_score = confidence_score + self.context.mark_task_completed(node.task, confidence_score) + yield emit(f"[POLICY] Subtask validated: '{node.task[:50]}...'") return - # If between 20%-60% elif decision == "expand" and depth < self.max_depth: - # Use UNIFIED context for planner (same as executor/router) - planner_context = f""" -{self.context.get_unified_context(max_tokens=5000)} - -## Current Plan Status -{self.context.get_tasks(include_goal=False)} - -## Instructions -Analyze what has been achieved and expand the plan with only what still needs to be done. -Update confidence_score based on progress. Reason step by step for the most logical plan. -""" + planner_context = ( + f"{self.context.get_unified_context(max_tokens=5000)}\n\n" + f"## Current Plan Status\n{self.context.get_tasks(include_goal=False)}\n\n" + "## Instructions\n" + "Analyze what has been achieved and expand the plan with only what still needs to be done.\n" + "Update confidence_score based on progress. Reason step by step for the most logical plan.\n" + ) subtasks, website_info, exploit_info = await self.planner.expand( node, context=planner_context, usage=RunUsage(), - usage_limits=UsageLimits(request_limit=None) + usage_limits=UsageLimits(request_limit=None), ) formatted_website_info = _format_dict_for_context(website_info.model_dump()) self.context.add_tool_response("website_info", formatted_website_info) if exploit_info.reasoning or exploit_info.highly_possible_vulnerabilities: formatted_exploit_info = _format_dict_for_context(exploit_info.model_dump()) self.context.add_tool_response("exploit_info", formatted_exploit_info) - planner_subtasks = [] - # Because it's not hashable - for subtask in subtasks: - planner_subtask = TaskPlanner( - task=subtask.task, - confidence_score=subtask.confidence_score, - status=subtask.status - ) - planner_subtasks.append(planner_subtask) + planner_subtasks = [ + TaskPlanner(task=s.task, confidence_score=s.confidence_score, status=s.status) + for s in subtasks + ] parent_planner = TaskPlanner( task=node.task, confidence_score=node.confidence_score, - status=node.status + status=node.status, ) self.context.add_tasks(parent_task=parent_planner, tasks=planner_subtasks) + if not subtasks: node.status = "refine" - yield emit(f"[PLANNER] No subtasks generated for '{node.task}', requesting \ - refinement") + yield emit(f"[PLANNER] No subtasks generated for '{node.task}', requesting refinement") if node.parent: - async for chunk in self._solve(node.parent, depth=node.depth, exit_strategy=exit_strategy, exit_loop=exit_loop): - # Check if child call signaled exit_loop - if isinstance(chunk, dict) and chunk.get('exit_loop'): + async for chunk in self._solve(node.parent, depth=node.depth, exit_loop=exit_loop): + if isinstance(chunk, dict) and chunk.get("exit_loop"): should_exit = True yield chunk break @@ -331,36 +355,30 @@ def emit(message: str) -> str: node.children = subtasks yield emit(f"[PLANNER] Generated {len(subtasks)} subtasks for '{node.task}'") for subtask in subtasks: - async for chunk in self._solve(subtask, depth + 1, exit_strategy=exit_strategy, exit_loop=exit_loop): - # Check if child call signaled exit_loop - if isinstance(chunk, dict) and chunk.get('exit_loop'): + async for chunk in self._solve(subtask, depth + 1, exit_loop=exit_loop): + if isinstance(chunk, dict) and chunk.get("exit_loop"): should_exit = True yield chunk break yield chunk else: - # Continue to next subtask if no exit_loop was signaled continue - # Break out of subtask loop if exit_loop was signaled break - # If refine - else: - # Use UNIFIED context for planner (same as executor/router) - planner_context = f""" -{self.context.get_unified_context(max_tokens=5000)} -## Current Plan Status -{self.context.get_tasks(include_goal=False)} - -## Instructions -Analyze what has been achieved and update the plan with only what still needs to be done. -Update confidence_score for completed items. Reason step by step for the most logical updated plan. -""" + # refine (60-80% or expand at max depth) + else: + planner_context = ( + f"{self.context.get_unified_context(max_tokens=5000)}\n\n" + f"## Current Plan Status\n{self.context.get_tasks(include_goal=False)}\n\n" + "## Instructions\n" + "Analyze what has been achieved and update the plan with only what still needs to be done.\n" + "Update confidence_score for completed items. Reason step by step for the most logical updated plan.\n" + ) updated_tasks, website_info, exploit_info = await self.planner.update_plan( node, context=planner_context, usage=RunUsage(), - usage_limits=UsageLimits(request_limit=None) + usage_limits=UsageLimits(request_limit=None), ) formatted_website_info = _format_dict_for_context(website_info.model_dump()) self.context.add_tool_response("website_info", formatted_website_info) @@ -368,53 +386,42 @@ def emit(message: str) -> str: formatted_exploit_info = _format_dict_for_context(exploit_info.model_dump()) self.context.add_tool_response("exploit_info", formatted_exploit_info) - # Store parent reference before updating node parent_task = node.parent - # Update the parent's children with the updated tasks if parent_task: parent_task.children = updated_tasks - # Update the current node to the updated version - # Try to find the updated version by matching the task description first - # If not found, use the first updated task (assuming it's the same task refined) updated_node = None for updated_task in updated_tasks: if updated_task.task == node.task: updated_node = updated_task break - # If exact match not found, use first task (task might have been refined/renamed) node = updated_node if updated_node else updated_tasks[0] if updated_tasks else node else: - # If no parent, update the node itself if updated_tasks: node = updated_tasks[0] - # Update context with updated tasks - planner_subtasks = [] - for subtask in updated_tasks: - planner_subtask = TaskPlanner( - task=subtask.task, - confidence_score=subtask.confidence_score, - status=subtask.status - ) - planner_subtasks.append(planner_subtask) + planner_subtasks = [ + TaskPlanner(task=s.task, confidence_score=s.confidence_score, status=s.status) + for s in updated_tasks + ] if parent_task: parent_planner = TaskPlanner( task=parent_task.task, confidence_score=parent_task.confidence_score, - status=parent_task.status + status=parent_task.status, ) self.context.add_tasks(parent_task=parent_planner, tasks=planner_subtasks) else: self.context.add_tasks(parent_task=None, tasks=planner_subtasks) - yield emit(f"[PLANNER] Updated plan for tasks \ - with parent '{parent_task.task if parent_task else 'root'}'") + yield emit( + f"[PLANNER] Updated plan for tasks " + f"with parent '{parent_task.task if parent_task else 'root'}'" + ) - async for chunk in self._solve(node=node, depth=depth, exit_strategy=exit_strategy, exit_loop=exit_loop): - # Check if child call signaled exit_loop - if isinstance(chunk, dict) and chunk.get('exit_loop'): + async for chunk in self._solve(node=node, depth=depth, exit_loop=exit_loop): + if isinstance(chunk, dict) and chunk.get("exit_loop"): should_exit = True yield chunk break @@ -422,18 +429,16 @@ def emit(message: str) -> str: def _policy(self, confidence_score: float) -> Literal["fail", "expand", "refine", "validate"]: - """ - The policy is given a confidence score, and depending on the information given - is capable to determine the next step for the task. - - <20% : treated as unrecoverable. Even though the goal is well defined, we need - to propagate the failure and move on. - - 20-60% : stays in exploration mode, this stage trigger more granular subtasks - or gather missing information. - - 60%-80% : keep executing and reiterating. At this stage we can either make simple changes - to increase the confidence score or run the generic agents for executing and testing. - - > 80% : move to validator/controller. - the policy could be assertions, tool verification, or LLM judge - and before the task is done. + """Determine the next action for a subtask based on its confidence score. + + - <20% : fail — unrecoverable, propagate failure. + - 20-60%: expand — decompose into finer-grained subtasks. + - 60-80%: refine — iterate on the same task. + - >=80% : validate — subtask confidence is high, mark done. + + Note: root-goal validation is handled by the ValidationGate, + not by this policy. 'validate' here only means the *subtask* + has high enough confidence to be considered complete. """ if confidence_score < self.FAIL_THRESHOLD: return "fail" @@ -443,85 +448,42 @@ def _policy(self, confidence_score: float) -> Literal["fail", "expand", "refine" return "validate" return "refine" - async def _validate(self, node: TaskNode) -> Tuple[str, str]: - """Validate a task node's execution. - - Args: - node: The TaskNode to validate - - Returns: - Status string: "completed" if validation passes, "failed-validation" otherwise. - If no validator is available, returns "completed" by default. + async def _write_report(self, verdict: ValidationVerdict) -> None: + """Delegate report generation and persistence to the ReporterAgent. - Note: - Updates the node's confidence_score with the validation confidence score. + The reporter agent has the write_workspace_file tool and writes the report + itself during execution — no programmatic write after the call. + Uses agent_id (not session_id) to match the mounted AVFS workspace. """ - if not self.validator: - return ("completed", "") - - # Use UNIFIED context for validator (same as executor/router/planner) - # This ensures validator sees the same discoveries and exploits as other agents - validation_context_text = self.context.get_unified_context(max_tokens=5000) - - (valid, confidence_score, critique, validation_token) = await self.validator.verify( - task=node, - context=validation_context_text - ) - - # Record validation result compactly - validation_summary = f"Validation: {critique[:100]}, confidence: {confidence_score:.2f}" - if validation_token: - validation_summary += f", token: {validation_token}" - - # Add to structured context only (avoid bloating workflow_context) - self.context.structured.append_to_log(validation_summary) - self.context.add_agent_response(validation_summary, skip_structured=True) - - node.confidence_score = confidence_score - - # If validation passed, add validated result as high-confidence fact - # This ensures the successful exploit details persist in context for next agents - if valid: - # Extract recent successful attempts to preserve as facts - successful_attempts = [a for a in self.context.structured.attempts if a.result == "success"] - for attempt in successful_attempts[-3:]: # Last 3 successful - self.context.structured.add_fact_simple( - category="validated_exploit", - key=f"{node.task[:50]}", - value=f"Payload: {attempt.payload}", - confidence=confidence_score, - source_task=node.task, - details={ - "payload": attempt.payload, - "reason": attempt.reason, - "validation_token": validation_token, - "task": attempt.task - }, - actionable=True - ) - - return ("completed", validation_token) if valid else ("failed-validation", validation_token) + try: + await self.reporter.summarize_and_write( + root_goal=self.context.final_goal, + verdict=verdict, + context=self.context.get_unified_context(max_tokens=100_000), + session_id=str(self.agent_id), + ) + except Exception as exc: + # Report writing must never crash the agent loop. + logger.warning("ReporterAgent failed to write report: %s", exc) async def run( self, task: str, - exit_strategy: str, context: str | None = None, ) -> AsyncGenerator[str | dict[str, Any], None]: """Run the ADaPT agent on a given task. - Creates a root task node and recursively solves it using the ADaPT algorithm, - which includes execution, planning, and validation phases. + Creates a root task node and recursively solves it using the ADaPT + algorithm, which includes execution, planning, and validation phases. Args: - task: The main task description to execute - exit_strategy: Strategy for determining when to exit + task: The main task description to execute. + context: Optional initial context string for the planner. Yields: - Human-readable log strings describing progress followed by a final - {"type": "result", "root": TaskNode} event containing the execution tree. + Human-readable log strings describing progress followed by a + final {"type": "result", "root": TaskNode} event. """ - # Reset attempt tracking for new run self._attempted_tasks.clear() root = TaskNode( @@ -530,22 +492,18 @@ async def run( confidence_score=0.7, status="pending", parent=None, - children=[] + children=[], ) self.context.set_root_task(root.task) - # # Use UNIFIED context for initial planning (same as all other agents) - # initial_context = self.context.get_unified_context(max_tokens=4000) subtasks, website_info, exploit_info = await self.planner.expand( root, context=context, usage=RunUsage(), - usage_limits=UsageLimits(request_limit=None) + usage_limits=UsageLimits(request_limit=None), ) - # print(f"subtasks: {subtasks}") - # print(f"website info: {website_info}") - # Add website info to structured context as facts (not verbose dump) + # Store website info as discovered facts. website_dict = website_info.model_dump() if website_dict.get("information_gathering"): info = website_dict["information_gathering"] @@ -556,56 +514,48 @@ async def run( category="website_info", key=key, value=str(value)[:200], - confidence=0.8 + confidence=0.8, ) else: self.context.add_discovered_fact( category="website_info", key="general", value=str(info)[:200], - confidence=0.8 + confidence=0.8, ) - # Add exploit info if available + # Store exploit info as discovered facts. if exploit_info.reasoning or exploit_info.highly_possible_vulnerabilities: if exploit_info.highly_possible_vulnerabilities: - # Handle both string and list formats vulns_raw = exploit_info.highly_possible_vulnerabilities if isinstance(vulns_raw, str): - # Split by comma or newline if it's a string - if ',' in vulns_raw: - vulns = [v.strip() for v in vulns_raw.split(',') if v.strip()] - elif '\n' in vulns_raw: - vulns = [v.strip() for v in vulns_raw.split('\n') if v.strip()] + if "," in vulns_raw: + vulns = [v.strip() for v in vulns_raw.split(",") if v.strip()] + elif "\n" in vulns_raw: + vulns = [v.strip() for v in vulns_raw.split("\n") if v.strip()] else: - # Single vulnerability as string vulns = [vulns_raw.strip()] if vulns_raw.strip() else [] else: vulns = list(vulns_raw) if vulns_raw else [] - for vuln in vulns[:5]: # Limit to 5 + for vuln in vulns[:5]: self.context.add_discovered_fact( category="vulnerability", key=str(vuln)[:50], value=str(vuln), - confidence=0.6 + confidence=0.6, ) - planner_subtasks = [] - for subtask in subtasks: - planner_subtask = TaskPlanner( - task=subtask.task, - confidence_score=subtask.confidence_score, - status=subtask.status - ) - planner_subtasks.append(planner_subtask) - # TODO: handling the termination + + planner_subtasks = [ + TaskPlanner(task=s.task, confidence_score=s.confidence_score, status=s.status) + for s in subtasks + ] self.context.add_tasks(parent_task=None, tasks=planner_subtasks) - # print(f"task context is \n {self.context.get_tasks(0)}") + exit_loop_triggered = False for subtask in subtasks: - async for chunk in self._solve(subtask, depth=1, exit_strategy=exit_strategy, exit_loop=False): - # Check if exit_loop was signaled - if isinstance(chunk, dict) and chunk.get('exit_loop'): + async for chunk in self._solve(subtask, depth=1, exit_loop=False): + if isinstance(chunk, dict) and chunk.get("exit_loop"): exit_loop_triggered = True yield chunk break @@ -613,8 +563,5 @@ async def run( if exit_loop_triggered: break - # Always yield the final result, even when exit_loop was triggered - # Update root children to include all subtasks root.children = subtasks yield {"type": "result", "root": root} - return diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py index 14ebd73..4d2d6b2 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py @@ -374,8 +374,6 @@ def _persist_agent_summary(agent_name: str, task: str, output: AgentOutput) -> N @supervisor.agent.tool async def call_requester_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: """Call the requester agent to perform HTTP request testing.""" - print(f"input tool looking for the error : {prompt}") - if ctx.deps.requester_agent is None or ctx.deps.requester_deps is None: return "Requester agent dependencies not configured." memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) @@ -388,7 +386,6 @@ async def call_requester_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> deferred_tool_results=ctx.deps.deferred_tool_results ) if hasattr(result, 'output') and isinstance(result.output, AgentOutput): - # Add output to context for future reference _add_agent_output_to_context( task=task_node.task, context=ctx.deps.context, @@ -396,13 +393,15 @@ async def call_requester_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> output=result.output ) _persist_agent_summary("requester", prompt, result.output) - return f"Requester agent result: {result.output.model_dump()}" - return str(result.output) if hasattr(result, 'output') else str(result) + result_str = f"Requester agent result: {result.output.model_dump()}" + else: + result_str = str(result.output) if hasattr(result, 'output') else str(result) + emit(f"[requester] prompt={prompt[:200]} | result={result_str[:300]}") + return result_str @supervisor.agent.tool async def call_shell_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: """Call the shell agent to execute shell commands.""" - print(f"input tool looking for the error : {prompt}") if ctx.deps.shell_agent is None or ctx.deps.shell_deps is None: return "Shell agent dependencies not configured." memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) @@ -415,7 +414,6 @@ async def call_shell_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: deferred_tool_results=ctx.deps.deferred_tool_results ) if hasattr(result, 'output') and isinstance(result.output, AgentOutput): - # Add output to context for future reference _add_agent_output_to_context( task=task_node.task, context=ctx.deps.context, @@ -423,14 +421,15 @@ async def call_shell_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: output=result.output ) _persist_agent_summary("shell", prompt, result.output) - return f"Shell agent result: {result.output.model_dump()}" - return str(result.output) if hasattr(result, 'output') else str(result) + result_str = f"Shell agent result: {result.output.model_dump()}" + else: + result_str = str(result.output) if hasattr(result, 'output') else str(result) + emit(f"[shell] prompt={prompt[:200]} | result={result_str[:300]}") + return result_str @supervisor.agent.tool async def call_webapp_analyzer_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: - print(f"input tool looking for the error : {prompt}") - - print(ctx.deps.requester_deps) + """Call the webapp analyzer agent to analyze web application structure and behavior.""" memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) result = await ctx.deps.webapp_analyzer_agent.run( f"{memory_prefix}{prompt}", @@ -438,16 +437,15 @@ async def call_webapp_analyzer_agent(ctx: RunContext[SupervisorDeps], prompt: st message_history=ctx.deps.message_history, usage=ctx.usage, usage_limits=ctx.deps.usage_limits, - deferred_tool_results=ctx.deps.deferred_tool_results + deferred_tool_results=ctx.deps.deferred_tool_results ) - - return str(result.output.model_dump()) + result_str = str(result.output.model_dump()) + emit(f"[webapp_analyzer] prompt={prompt[:200]} | result={result_str[:300]}") + return result_str @supervisor.agent.tool async def call_python_interpreter_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: """Call the python interpreter agent to execute Python scripts.""" - print(f"input tool looking for the error : {prompt}") - memory_prefix = _memory_prompt_prefix(ctx.deps.memory_context) result = await ctx.deps.python_interpreter_agent.run( f"{memory_prefix}{prompt}", @@ -459,7 +457,6 @@ async def call_python_interpreter_agent(ctx: RunContext[SupervisorDeps], prompt: deferred_tool_results=ctx.deps.deferred_tool_results ) if hasattr(result, 'output') and isinstance(result.output, AgentOutput): - # Add output to context for future reference _add_agent_output_to_context( task=task_node.task, context=ctx.deps.context, @@ -467,8 +464,11 @@ async def call_python_interpreter_agent(ctx: RunContext[SupervisorDeps], prompt: output=result.output ) _persist_agent_summary("python_interpreter", prompt, result.output) - return f"Python interpreter agent result: {result.output.model_dump()}" - return str(result.output) if hasattr(result, 'output') else str(result) + result_str = f"Python interpreter agent result: {result.output.model_dump()}" + else: + result_str = str(result.output) if hasattr(result, 'output') else str(result) + emit(f"[python_interpreter] prompt={prompt[:200]} | result={result_str[:300]}") + return result_str @supervisor.agent.tool async def call_memory_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str: diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py new file mode 100644 index 0000000..c47d463 --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py @@ -0,0 +1,455 @@ +# Copyright (C) 2025 Yassine Bargach +# Licensed under the GNU Affero General Public License v3 +# See LICENSE file for full license information. + +"""Validation strategies for the ADaPT agent architecture. + +This module defines a composable validation system that determines whether +the root goal of a security assessment has been achieved. Strategies are +chained in a ValidationGate — the gate iterates each strategy in order +and short-circuits on the first definitive stop signal. + +Configuration is loaded from a YAML file (default: +``~/.cache/deadend/validation.yaml``). See ``ValidationConfig`` for the +schema and ``load_validation_config`` for the loader. + +Adding a new strategy: + 1. Create a class that implements the ValidationStrategy protocol. + 2. Register it in STRATEGY_REGISTRY at the bottom of this file. + 3. Reference it by name in the YAML ``strategies`` list. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Protocol, runtime_checkable + +import yaml +from pydantic import BaseModel, Field + +from deadend_agent.config.settings import ModelSpec +from deadend_agent.logging import logger +from deadend_prompts import render_agent_instructions + + +# --------------------------------------------------------------------------- +# Default config path +# --------------------------------------------------------------------------- + +DEFAULT_CONFIG_PATH = Path.home() / ".cache" / "deadend" / "validation.yaml" + + +# --------------------------------------------------------------------------- +# Data structures +# --------------------------------------------------------------------------- + +class ValidationVerdict(BaseModel): + """Result of a validation strategy check. + + Attributes: + stop: True if the root goal is satisfied — the agent should exit. + confidence: How confident the strategy is in this verdict (0.0–1.0). + token: Captured validation token (flag, proof string, etc.). + critique: Human-readable explanation of the verdict. + report: Markdown report content ready to be written to disk. + """ + stop: bool = False + confidence: float = 0.0 + token: str = "" + critique: str = "" + report: str = "" + + +# --------------------------------------------------------------------------- +# YAML-driven configuration +# --------------------------------------------------------------------------- + +class StrategyConfig(BaseModel): + """Per-strategy configuration block in the YAML file. + + Example YAML:: + + strategies: + - name: flag + pattern: "FLAG\\{[^}]+\\}" + - name: judge + validation_type: flag + validation_format: "FLAG{}" + """ + name: str + pattern: str | None = None + validation_type: str | None = None + validation_format: str | None = None + + +class ValidationConfig(BaseModel): + """Top-level validation configuration loaded from YAML. + + Attributes: + strategies: Ordered list of strategy configs to chain. + validation_format: Global format string injected into agent prompts. + Individual strategy-level ``validation_format`` overrides this. + validation_type: Global type string injected into agent prompts. + """ + strategies: list[StrategyConfig] = Field( + default_factory=lambda: [ + StrategyConfig(name="flag"), + StrategyConfig(name="judge"), + ] + ) + validation_format: str | None = None + validation_type: str | None = None + + +def load_validation_config( + path: str | Path | None = None, +) -> ValidationConfig: + """Load a ``ValidationConfig`` from a YAML file. + + If *path* is ``None``, falls back to ``DEFAULT_CONFIG_PATH``. + If the file does not exist, returns the default config (flag + judge). + + Args: + path: Filesystem path to the YAML file. + + Returns: + A parsed and validated ``ValidationConfig``. + """ + config_path = Path(path) if path else DEFAULT_CONFIG_PATH + + if not config_path.exists(): + logger.debug( + "Validation config not found at %s — using defaults.", config_path, + ) + return ValidationConfig() + + try: + raw = yaml.safe_load(config_path.read_text(encoding="utf-8")) + except (yaml.YAMLError, OSError) as exc: + logger.warning( + "Failed to parse validation config at %s: %s — using defaults.", + config_path, exc, + ) + return ValidationConfig() + + if not isinstance(raw, dict): + logger.warning("Validation config is not a YAML mapping — using defaults.") + return ValidationConfig() + + return ValidationConfig(**raw) + + +# --------------------------------------------------------------------------- +# Strategy protocol +# --------------------------------------------------------------------------- + +@runtime_checkable +class ValidationStrategy(Protocol): + """Interface that every validation strategy must implement.""" + + async def check( + self, + output: dict, + root_goal: str, + context: str, + ) -> ValidationVerdict: + """Evaluate whether *root_goal* is satisfied. + + Args: + output: The supervisor output dict (keys: task_achieved, + detailed_summary, proofs, confidence_score). + root_goal: The top-level goal of the entire assessment. + context: Accumulated execution context (unified context string). + + Returns: + A ValidationVerdict indicating whether to stop. + """ + ... + + +# --------------------------------------------------------------------------- +# FlagStrategy — deterministic regex, zero LLM cost +# --------------------------------------------------------------------------- + +class FlagStrategy: + """Scan supervisor output and context for a flag matching a regex pattern. + + This is a pure string-search strategy — no LLM call, no network call. + It is always safe to run on every supervisor return. + """ + + def __init__(self, pattern: str = r"FLAG\{[^}]+\}"): + self.pattern = re.compile(pattern, re.IGNORECASE) + + async def check( + self, + output: dict, + root_goal: str, + context: str, + ) -> ValidationVerdict: + # Search in order: proofs first (most likely), then summary, then full context. + searchable_fields = [ + output.get("proofs", ""), + output.get("detailed_summary", ""), + context, + ] + for field in searchable_fields: + match = self.pattern.search(field) + if match: + token = match.group(0) + logger.debug("FlagStrategy: matched token '%s'", token) + return ValidationVerdict( + stop=True, + confidence=1.0, + token=token, + critique=f"Flag found via pattern match: {token}", + report=self._build_report(output, token, root_goal), + ) + + return ValidationVerdict(stop=False, confidence=0.0) + + @staticmethod + def _build_report(output: dict, token: str, root_goal: str) -> str: + return ( + "# Validation Report — Flag Captured\n\n" + f"**Goal:** {root_goal}\n\n" + f"**Token:** `{token}`\n\n" + "## Evidence\n\n" + f"```\n{output.get('proofs', 'N/A')}\n```\n\n" + "## Summary\n\n" + f"{output.get('detailed_summary', 'N/A')}\n" + ) + + +# --------------------------------------------------------------------------- +# JudgeAgentStrategy — LLM agent that can re-verify findings +# --------------------------------------------------------------------------- + +class JudgeAgentStrategy: + """LLM-based judge that evaluates whether the root goal is satisfied. + + Unlike a passive judge, this is backed by an AgentRunner and can be + extended with tools (requester, shell, memory) to actively re-verify + findings. It self-throttles: if no new evidence has appeared since the + last check it returns a non-stop verdict immediately. + """ + + def __init__( + self, + model: ModelSpec, + validation_type: str | None = None, + validation_format: str | None = None, + ): + from deadend_agent.agents.factory import AgentRunner + + judge_instructions = render_agent_instructions( + "judge", + tools={}, + validation_type=validation_type or "flag", + validation_format=validation_format or "FLAG{{}}", + ) + + self._output_model = _JudgeOutput + self._agent = AgentRunner( + name="judge_validator", + model=model, + instructions=judge_instructions, + deps_type=None, + output_type=_JudgeOutput, + tools=[], + ) + + # Self-throttle: track context hash to avoid redundant calls. + self._last_context_hash: int = 0 + + async def check( + self, + output: dict, + root_goal: str, + context: str, + ) -> ValidationVerdict: + # Self-throttle: skip if context hasn't changed since last check. + current_hash = hash(context) + if current_hash == self._last_context_hash: + logger.debug("JudgeAgentStrategy: no new evidence, skipping.") + return ValidationVerdict(stop=False, confidence=0.0) + self._last_context_hash = current_hash + + prompt = ( + "# Objective\n" + "You are the Validator. Judge whether the following goal is " + "satisfied based on the execution trace.\n\n" + f"## Goal\n{root_goal}\n\n" + f"## Latest Supervisor Output\n" + f"Summary: {output.get('detailed_summary', '')}\n" + f"Proofs: {output.get('proofs', '')}\n" + f"Confidence: {output.get('confidence_score', 0.0)}\n\n" + "# Execution Trace\n" + f"{context}\n\n" + "# Instructions\n" + "- valid (true/false)\n" + "- confidence (float 0.00–1.00)\n" + "- critique (string)\n" + "- validation_token — copy the exact FLAG{...} if present, " + "otherwise empty string\n" + ) + + result = await self._agent.run( + prompt=prompt, + deps=None, + message_history="", + usage=None, + usage_limits=None, + deferred_tool_results=None, + ) + + judge_output = result.output + if not isinstance(judge_output, _JudgeOutput): + logger.debug("JudgeAgentStrategy: unexpected output type %s", type(judge_output)) + return ValidationVerdict(stop=False, confidence=0.0) + print(judge_output) + return ValidationVerdict( + stop=judge_output.valid, + confidence=judge_output.confidence_score, + token=judge_output.validation_token or "", + critique=judge_output.critique, + report=self._build_report(judge_output, root_goal) if judge_output.valid else "", + ) + + @staticmethod + def _build_report(judge: _JudgeOutput, root_goal: str) -> str: + token_line = f"**Token:** `{judge.validation_token}`\n\n" if judge.validation_token else "" + return ( + "# Validation Report — Judge Verdict\n\n" + f"**Goal:** {root_goal}\n\n" + f"**Verdict:** {'ACHIEVED' if judge.valid else 'NOT ACHIEVED'}\n" + f"**Confidence:** {judge.confidence_score:.2f}\n\n" + f"{token_line}" + "## Critique\n\n" + f"{judge.critique}\n" + ) + + +class _JudgeOutput(BaseModel): + """Structured output expected from the judge LLM call.""" + valid: bool + confidence_score: float + critique: str + validation_token: str | None = None + + +# --------------------------------------------------------------------------- +# ValidationGate — composite that chains strategies +# --------------------------------------------------------------------------- + +class ValidationGate: + """Runs a chain of validation strategies and short-circuits on the first stop. + + Usage:: + + gate = ValidationGate([FlagStrategy(), JudgeAgentStrategy(model)]) + verdict = await gate.check(output, root_goal, context) + if verdict.stop: + ... # write report, exit loop + """ + + def __init__(self, strategies: list[ValidationStrategy]): + if not strategies: + raise ValueError("ValidationGate requires at least one strategy.") + self.strategies = strategies + + async def check( + self, + output: dict, + root_goal: str, + context: str, + ) -> ValidationVerdict: + """Iterate strategies in order; return first stop=True verdict.""" + last_verdict = ValidationVerdict(stop=False, confidence=0.0) + + for strategy in self.strategies: + verdict = await strategy.check(output, root_goal, context) + if verdict.stop: + logger.debug( + "ValidationGate: stop signalled by %s (confidence=%.2f)", + type(strategy).__name__, + verdict.confidence, + ) + return verdict + # Keep the highest-confidence non-stop verdict for callers. + if verdict.confidence > last_verdict.confidence: + last_verdict = verdict + + return last_verdict + + +# --------------------------------------------------------------------------- +# Strategy registry & factory +# --------------------------------------------------------------------------- + +STRATEGY_REGISTRY: dict[str, type] = { + "flag": FlagStrategy, + "judge": JudgeAgentStrategy, +} + +# Presets: ordered lists of strategy names. +PRESETS: dict[str, list[str]] = { + "flag": ["flag"], + "judge": ["judge"], + "ctf": ["flag", "judge"], + "recon": ["judge"], +} + + +def build_validation_gate( + *, + config: ValidationConfig | None = None, + model: ModelSpec | None = None, +) -> ValidationGate: + """Build a ValidationGate from a ``ValidationConfig``. + + The config is typically loaded from ``validation.yaml`` via + ``load_validation_config()``. If *config* is ``None`` the default + config is used (flag + judge). + + Args: + config: Parsed YAML config. ``None`` → default. + model: Required when any LLM-based strategy is in the chain. + + Returns: + A configured ``ValidationGate`` ready for use. + """ + cfg = config or ValidationConfig() + + instances: list[ValidationStrategy] = [] + for strategy_cfg in cfg.strategies: + name = strategy_cfg.name + + if name == "flag": + pattern = strategy_cfg.pattern or r"FLAG\{[^}]+\}" + instances.append(FlagStrategy(pattern=pattern)) + + elif name == "judge": + if model is None: + raise ValueError("JudgeAgentStrategy requires a model.") + instances.append(JudgeAgentStrategy( + model=model, + validation_type=( + strategy_cfg.validation_type or cfg.validation_type + ), + validation_format=( + strategy_cfg.validation_format or cfg.validation_format + ), + )) + + elif name in STRATEGY_REGISTRY: + instances.append(STRATEGY_REGISTRY[name]()) # type: ignore[call-arg] + + else: + raise ValueError( + f"Unknown strategy '{name}'. " + f"Available: {list(STRATEGY_REGISTRY.keys())}" + ) + + return ValidationGate(instances) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validator.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validator.py deleted file mode 100644 index 9fd3adf..0000000 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validator.py +++ /dev/null @@ -1,123 +0,0 @@ -from pydantic import BaseModel -from deadend_agent.agents.components.planner import TaskNode -from deadend_agent.config.settings import ModelSpec -from deadend_agent.agents import AgentRunner -from deadend_prompts import render_agent_instructions - - -class ValidatorOutput(BaseModel): - """Output format for task validation results. - - Attributes: - valid: Boolean indicating whether the task execution is valid - confidence_score: Confidence score (0.0-1.0) for the validation decision - critique: Explanation or critique of the validation decision - """ - valid: bool - confidence_score: float - critique: str - validation_token: str | None - -class Validator: - """Validator component that double-checks task execution results. - - The validator uses an LLM-based judge method to verify that task executions - are coherent, valid, and meet the intended objectives. It provides validation - decisions along with confidence scores and critiques. - """ - - def __init__( - self, - model: ModelSpec, - validation_type: str | None, - validation_format: str | None, - ) -> None: - """Initialize the Validator. - - Args: - model: The AI model to use for validation/judgment - """ - if validation_type and validation_format: - judge_instructions = render_agent_instructions( - "judge", - tools={}, - validation_type=validation_type, - validation_format=validation_format - ) - else: - judge_instructions = render_agent_instructions( - "judge", - tools={}, - validation_type="flag", - validation_format="FLAG{{}}" - ) - self.agent = AgentRunner( - name="validator", - model=model, - deps_type=None, - instructions=judge_instructions, - output_type=ValidatorOutput, - tools=[] - ) - - async def verify(self, task: TaskNode, context: str) -> tuple[bool, float, str, str]: - """Verify whether a task execution is valid and successful. - - Args: - task: The TaskNode representing the task to validate - context: Execution context containing logs and execution trace - - Returns: - Tuple of (valid, confidence_score, critique) where: - - valid: Boolean indicating if the task execution is valid - - confidence_score: Float (0.0-1.0) indicating validation confidence - - critique: String explanation of the validation decision - """ - # Extract root goal from context (first line typically contains it) - # Parse the goal from unified context format "Goal: ..." - root_goal = "" - for line in context.split("\n"): - if line.startswith("Goal:"): - root_goal = line.replace("Goal:", "").strip() - break - - # If we found a root goal, validate against that. Otherwise use task.task - validation_objective = root_goal if root_goal else task.task - - prompt = f"""\ -# Objective -You are the Validator. Judge whether the following goal is satisfied based on the execution trace. -{validation_objective} - -# Execution trace -{context} - -# Output results -- valid (true/false) -- confidence (float 0.00-1.00) - like a percentage -- critique (string) -- validation_token corresponding to the result if found (only return if found). Copy the exact FLAG{{...}} if present. -""" - result = await self.agent.run( - prompt=prompt, - deps=None, - message_history="", - usage=None, - usage_limits=None, - deferred_tool_results=None - ) - - # returned variables initialization - valid = False - confidence_score = 0.0 - critique = "" - validation_token = "" - if isinstance(result.output, ValidatorOutput): - valid = result.output.valid - confidence_score = float(result.output.confidence_score) - critique = result.output.critique - validation_token = result.output.validation_token \ - if result.output.validation_token else "" - - # adding to context - return (valid, confidence_score, critique, validation_token) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py index 6fdc771..4d75873 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/generic_agents/memory_agent.py @@ -5,7 +5,12 @@ from deadend_agent.agents.factory import AgentRunner from deadend_agent.config.settings import ModelSpec -from deadend_agent.tools import avfs_grep, avfs_list, avfs_read, avfs_write +from deadend_agent.tools import ( + grep_memory_files, + list_memory_files, + read_memory_file, + write_memory_file, +) from deadend_prompts import render_agent_instructions, render_tool_description @@ -18,10 +23,10 @@ def __init__( deps_type: Any | None, ): tools_metadata = { - "avfs_list": render_tool_description("avfs_list"), - "avfs_read": render_tool_description("avfs_read"), - "avfs_write": render_tool_description("avfs_write"), - "avfs_grep": render_tool_description("avfs_grep"), + "list_memory_files": render_tool_description("list_memory_files"), + "read_memory_file": render_tool_description("read_memory_file"), + "write_memory_file": render_tool_description("write_memory_file"), + "grep_memory_files": render_tool_description("grep_memory_files"), } self.instructions = render_agent_instructions( @@ -36,10 +41,10 @@ def __init__( deps_type=deps_type, output_type=[str, DeferredToolRequests], tools=[ - Tool(avfs_list), - Tool(avfs_read), - Tool(avfs_write), - Tool(avfs_grep), + Tool(list_memory_files), + Tool(read_memory_file), + Tool(write_memory_file), + Tool(grep_memory_files), ], ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py index 567402a..0dffa4e 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/reporter.py @@ -2,41 +2,84 @@ # Licensed under the GNU Affero General Public License v3 # See LICENSE file for full license information. -"""Reporter agent for summarizing workflow context and maintaining token limits. +"""Reporter agent for summarizing assessment findings and writing reports. -This module implements an AI agent that analyzes the accumulated workflow context, -summarizes key information, and ensures the context remains within manageable token -limits (under 150,000 tokens) for optimal AI model performance and cost efficiency. +This agent has access to the AVFS write tool. When invoked it analyzes +the execution context, produces a structured markdown report, and +**writes it to disk itself** via the tool — there is no programmatic +write after the LLM call. """ -from pydantic import BaseModel +from __future__ import annotations + +from dataclasses import dataclass + +from pydantic_ai import Tool +from pydantic_ai.usage import RunUsage, UsageLimits + +from deadend_agent.agents.components.validation_strategies import ValidationVerdict +from deadend_agent.config.settings import ModelSpec from deadend_agent.context.context_engine import ContextEngine -from deadend_prompts import render_agent_instructions +from deadend_agent.logging import logger +from deadend_agent.tools import write_workspace_file +from deadend_prompts import render_agent_instructions, render_tool_description + from .factory import AgentRunner -class ReporterOutput(BaseModel): - summarized_context: str + +# --------------------------------------------------------------------------- +# Deps — the reporter only needs a session_id so write_workspace_file can resolve paths +# --------------------------------------------------------------------------- + +@dataclass +class ReporterDeps: + """Minimal deps for the reporter agent.""" + session_id: str + + +# --------------------------------------------------------------------------- +# Agent +# --------------------------------------------------------------------------- class ReporterAgent(AgentRunner): + """Reporter agent — summarizes findings and writes reports to AVFS. + + Unlike a passive summarizer, this agent has the ``write_workspace_file`` tool + and is instructed to write the report file itself during execution. + ``output_type`` is ``str`` (plain text confirmation) to avoid the + Azure AI ``tool_choice`` object format incompatibility that occurs + with pydantic output schemas. """ - Reporter Agent - """ - def __init__(self, model, deps_type, tools, validation_type: str | None, validation_format: str | None): + + def __init__( + self, + model: ModelSpec, + validation_type: str | None = None, + validation_format: str | None = None, + ): + tools_metadata = { + "write_workspace_file": render_tool_description("write_workspace_file"), + } + reporter_instructions = render_agent_instructions( - "reporter", - tools={}, - validation_type=validation_type, - validation_format=validation_format + "reporter", + tools=tools_metadata, + validation_type=validation_type or "security assessment", + validation_format=validation_format or "Information", ) - self._set_description() + super().__init__( name="reporter", model=model, instructions=reporter_instructions, - deps_type=deps_type, - output_type=ReporterOutput, - tools=[] + deps_type=ReporterDeps, + output_type=str, + tools=[Tool(write_workspace_file)], ) + self.description = ( + "The reporter summarizes assessment findings and writes reports." + ) + async def run( self, prompt, @@ -46,7 +89,7 @@ async def run( usage_limits, deferred_tool_results=None, *args, - **kwargs + **kwargs, ): return await super().run( prompt=prompt, @@ -54,48 +97,108 @@ async def run( message_history=message_history, usage=usage, usage_limits=usage_limits, - deferred_tool_results=None, + deferred_tool_results=deferred_tool_results, ) - async def summarize_context(self, context_engine: ContextEngine): - """Summarize the workflow context and update it using the context engine setter. - + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + async def summarize_and_write( + self, + root_goal: str, + verdict: ValidationVerdict, + context: str, + session_id: str, + ) -> str: + """Run the reporter agent to generate and write a report. + + The agent itself calls ``write_workspace_file`` to persist the report. + We do NOT write programmatically after the call — the agent + does it. + Args: - context_engine: The ContextEngine instance containing the workflow context to summarize. - + root_goal: The top-level objective of the assessment. + verdict: The ValidationVerdict that triggered the report. + context: Full accumulated execution context. + session_id: AVFS session identifier (passed via deps). + Returns: - ReporterOutput: The summarized context output. + The agent's text output (confirmation message). """ - # Get the current workflow context (with possible summarization) - current_context = context_engine.workflow_context + prompt = self._build_prompt(root_goal, verdict, context) + deps = ReporterDeps(session_id=session_id) + + result = await self.run( + prompt=prompt, + deps=deps, + message_history="", + usage=RunUsage(), + usage_limits=UsageLimits(request_limit=None), + ) + + output = str(result.output) if hasattr(result, "output") else str(result) + logger.debug("ReporterAgent output: %s", output[:200]) + return output - # Create a prompt for summarization - summarization_prompt = f""" -analyze and summarize the following workflow context while preserving all critical security information, vulnerabilities, and technical details. Keep the summary under 150,000 tokens while maintaining actionable intelligence for continued security testing. -Current workflow context: -{current_context} + async def summarize_context(self, context_engine: ContextEngine, session_id: str) -> str: + """Summarize the workflow context and write it. + + Legacy entry point used outside of the validation gate. """ + current_context = context_engine.workflow_context + prompt = ( + "Analyze and summarize the following workflow context into a " + "security assessment report. Preserve all critical security " + "information, vulnerabilities, and technical details.\n\n" + "Write the report to `reports/context_summary.md` using the " + "write_workspace_file tool.\n\n" + f"## Workflow Context\n{current_context}" + ) - # Run the reporter agent to get the summary + deps = ReporterDeps(session_id=session_id) result = await self.run( - prompt=summarization_prompt, - deps=None, + prompt=prompt, + deps=deps, message_history="", - usage=None, - usage_limits=None + usage=RunUsage(), + usage_limits=UsageLimits(request_limit=None), + ) + + output = str(result.output) if hasattr(result, "output") else str(result) + return output + + # ------------------------------------------------------------------ + # Prompt builder + # ------------------------------------------------------------------ + + @staticmethod + def _build_prompt( + root_goal: str, + verdict: ValidationVerdict, + context: str, + ) -> str: + token_line = ( + f"- Validation token found: `{verdict.token}`\n" + if verdict.token else "" + ) + return ( + "You are writing a security assessment report. Analyze the data " + "below and produce a comprehensive markdown report.\n\n" + "**You MUST write the report to `reports/validation_report.md` " + "using the write_workspace_file tool.** Do not just return the report as " + "text — call the tool to persist it.\n\n" + "IMPORTANT:\n" + "- Preserve EXACT working payloads character-for-character\n" + "- Include full HTTP requests that succeeded\n" + "- Include response snippets proving vulnerabilities\n" + "- Document filter bypass techniques with exact encoding used\n" + "- Note validation status (reflected vs executed, needs browser test)\n\n" + f"## Goal\n{root_goal}\n\n" + f"## Validation Result\n" + f"- Verdict: {'ACHIEVED' if verdict.stop else 'NOT ACHIEVED'}\n" + f"- Confidence: {verdict.confidence:.2f}\n" + f"{token_line}" + f"- Critique: {verdict.critique}\n\n" + f"## Assessment Data\n{context}\n" ) - # Extract the summarized context from the result - if hasattr(result, 'output') and hasattr(result.output, 'summarized_context'): - summarized_context = result.output.summarized_context - else: - # Fallback if the output structure is different - summarized_context = str(result.output) - print(f"result reporter : {result}") - print(f"result reporter output : {result.output}") - # Update the context engine with the summarized context - context_engine.set_new_workflow(summarized_context) - - return result - - def _set_description(self): - self.description = "The reporter summarize the context as it understood it." diff --git a/deadend_cli/deadend_agent/src/deadend_agent/config/validation.default.yaml b/deadend_cli/deadend_agent/src/deadend_agent/config/validation.default.yaml new file mode 100644 index 0000000..0175b6e --- /dev/null +++ b/deadend_cli/deadend_agent/src/deadend_agent/config/validation.default.yaml @@ -0,0 +1,62 @@ +# Validation Strategy Configuration +# ---------------------------------- +# This file defines how the agent validates whether the root goal +# has been achieved. Copy this to ~/.cache/deadend/validation.yaml +# to customize. +# +# Available strategies (chained in order, first stop=True wins): +# +# flag — Deterministic regex match on proofs/summary/context. +# Zero cost, runs every supervisor return. +# Config: pattern (regex string) +# +# judge — LLM agent that evaluates the execution trace against +# the root goal. Can re-verify findings. Self-throttles +# when no new evidence has appeared. +# Config: validation_type, validation_format +# +# Top-level validation_format is injected into all agent prompts +# (anti-fabrication guidance adapts to show token-specific rules). + +# --- CTF preset (default) --- +# Flag regex first (free), LLM judge as fallback. + +validation_format: "FLAG{}" +validation_type: "flag" + +strategies: + - name: flag + pattern: "FLAG\\{[^}]+\\}" + + - name: judge + +# --- Examples for other use cases --- +# +# # Recon / generic security assessment (no flag to match): +# validation_format: null +# validation_type: "security assessment" +# strategies: +# - name: judge +# +# # HackTheBox: +# validation_format: "HTB{}" +# validation_type: "flag" +# strategies: +# - name: flag +# pattern: "HTB\\{[^}]+\\}" +# - name: judge +# validation_format: "HTB{}" +# +# # picoCTF: +# validation_format: "picoCTF{}" +# validation_type: "flag" +# strategies: +# - name: flag +# pattern: "picoCTF\\{[^}]+\\}" +# - name: judge +# validation_format: "picoCTF{}" +# +# # Flag-only (no LLM judge, fastest): +# validation_format: "FLAG{}" +# strategies: +# - name: flag diff --git a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py index f428468..dbcd9e2 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/core_agent/core_agent.py @@ -205,11 +205,16 @@ def __init__( # We always *attempt* to use Instructor when available and an output_schema # is provided, but will gracefully fall back to manual JSON extraction # if the Instructor call fails for any reason. - if INSTRUCTOR_AVAILABLE and output_schema: + # Only use structured output for actual Pydantic BaseModel subclasses, + # not plain types like str, int, etc. + _is_pydantic_schema = ( + isinstance(output_schema, type) and issubclass(output_schema, BaseModel) + ) + if INSTRUCTOR_AVAILABLE and output_schema and _is_pydantic_schema: self.instructor_client = instructor.from_litellm(acompletion) else: self.instructor_client = None - if output_schema and not INSTRUCTOR_AVAILABLE: + if output_schema and not INSTRUCTOR_AVAILABLE and _is_pydantic_schema: logger.warning("instructor not available, structured output disabled") def tool(self, func: Callable) -> Callable: @@ -526,8 +531,14 @@ async def _run_impl( if usage_limits and self.tool_call_count >= usage_limits.get("tools", float('inf')): raise UsageLimitExceeded(f"Tool call limit reached: {usage_limits['tools']}") - # Extract structured output - if self.output_schema and self.instructor_client: + # Extract structured output when output_schema is a Pydantic BaseModel. + # Uses Instructor if available, otherwise falls back to manual JSON extraction. + _has_pydantic_schema = ( + self.output_schema + and isinstance(self.output_schema, type) + and issubclass(self.output_schema, BaseModel) + ) + if _has_pydantic_schema: output = await self._extract_structured(messages) else: # Return last assistant message content @@ -1237,6 +1248,7 @@ async def _extract_structured_manual(self, messages: list[dict]) -> BaseModel: kwargs = { "model": self.model, "messages": extraction_messages, + "response_format": {"type": "json_object"}, } if self.api_base: kwargs["api_base"] = self.api_base diff --git a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py index da7cdc3..b76e172 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py @@ -1,5 +1,4 @@ """Main DeadEnd agent orchestration module.""" -import re from pathlib import Path from typing import Any, Awaitable, Callable, Dict, Generator from uuid import UUID @@ -13,12 +12,17 @@ from deadend_agent.context import ContextEngine from deadend_agent.rag.sqlite_connector import SqliteRagConnector from deadend_agent.sandbox.sandbox import Sandbox -from deadend_agent.agents.reporter import ReporterAgent +from deadend_agent.agents.reporter import ReporterAgent, ReporterDeps from deadend_agent.agents.architecture import ADaPTAgent from deadend_agent.agents.generic_agents.memory_agent import MemoryAgent from deadend_agent.agents.components.executor import AgentExecutor, ResultEvent from deadend_agent.agents.components.planner import Planner, TaskNode -from deadend_agent.agents.components.validator import Validator +from deadend_agent.agents.components.validation_strategies import ( + ValidationConfig, + ValidationGate, + build_validation_gate, + load_validation_config, +) from deadend_agent.utils.structures import ( MemoryWorkspaceDeps, RequesterDeps, @@ -30,7 +34,6 @@ from deadend_agent.tools.avfs import avfs from .agents.recon_threatmodel_agent import ReconThreatModelAgent from .agents.exploit_web_agent import PlannerExploitAgent -# from deadend_eval.metrics imporst save_traces ApprovalCallback = Callable[..., Awaitable[str]] @@ -54,7 +57,8 @@ class DeadEndAgent: exploit_agent: PlannerExploitAgent | None = None planner: Planner executor: AgentExecutor - validator: Validator + validation_gate: ValidationGate + reporter: ReporterAgent adapt_agent: ADaPTAgent shell_deps: ShellDeps | None = None requester_deps: RequesterDeps | None = None @@ -69,8 +73,7 @@ def __init__( model: ModelSpec, available_agents: Dict[str, str], max_depth: int = 3, - validation_type: str | None = None, - validation_format: str | None = None, + validation_config_path: str | None = None, embedding_session_id: UUID | None = None, workspace_root: str | None = None, agents_storage_root: str | None = None, @@ -81,11 +84,23 @@ def __init__( self.max_depth = max_depth self.model = model self.available_agents = available_agents - self.validator = Validator( + + # Load validation config from YAML (falls back to defaults). + self.validation_config = load_validation_config(validation_config_path) + + # Build composable validation gate from the loaded config. + self.validation_gate = build_validation_gate( + config=self.validation_config, model=model, - validation_type=validation_type, - validation_format=validation_format ) + + # Reporter agent for writing assessment reports on validation stop. + self.reporter = ReporterAgent( + model=model, + validation_type=self.validation_config.validation_type, + validation_format=self.validation_config.validation_format, + ) + self.context = ContextEngine(model=self.model, session_id=session_id) self.workspace_root: str | None = None self.local_agent_id = local_agent_id or Config.get_local_agent_id() @@ -402,14 +417,13 @@ async def threat_model(self, task: str): reporter_agent = ReporterAgent( model=self.model, - deps_type=None, - tools=None, validation_format="Information", - validation_type="security assessment" + validation_type="security assessment", ) # context_text = await self.context.get_all_context() prompt_assessment = f"""\ Summarize the security assessment results from the reconnaissance phase. +Write the report to `reports/recon_report.md` using the write_workspace_file tool. IMPORTANT: - Preserve EXACT working payloads character-for-character @@ -423,7 +437,7 @@ async def threat_model(self, task: str): """ threat_model_data = await reporter_agent.run( prompt=prompt_assessment, - deps=None, + deps=ReporterDeps(session_id=str(self.agent_id)), usage=RunUsage(), usage_limits=UsageLimits(), deferred_tool_results=None, @@ -498,14 +512,13 @@ async def threat_model_stream(self, task: str): reporter_agent = ReporterAgent( model=self.model, - deps_type=None, - tools=None, validation_format="Information", - validation_type="security assessment" + validation_type="security assessment", ) # context_text = await self.context.get_all_context() prompt_assessment = f"""\ Summarize the security assessment results from the reconnaissance phase. +Write the report to `reports/recon_report.md` using the write_workspace_file tool. IMPORTANT: - Preserve EXACT working payloads character-for-character @@ -519,7 +532,7 @@ async def threat_model_stream(self, task: str): """ threat_model_data = await reporter_agent.run( prompt=prompt_assessment, - deps=None, + deps=ReporterDeps(session_id=str(self.agent_id)), usage=RunUsage(), usage_limits=UsageLimits(), deferred_tool_results=None, @@ -549,11 +562,13 @@ async def run_exploitation(self, threat_model: str, task: str): self.planner = Planner(planner_agent=self.exploit_agent, deps=self._target_session_key()) self.adapt_agent = ADaPTAgent( session_id=self.session_id, + agent_id=self.agent_id, context=self.context, executor=self.executor, planner=self.planner, - validator=self.validator, - max_depth=self.max_depth + validation_gate=self.validation_gate, + reporter=self.reporter, + max_depth=self.max_depth, ) plan: TaskNode | None = None @@ -575,7 +590,7 @@ async def run_exploitation(self, threat_model: str, task: str): traces: list[str | dict[str, Any]] = [] validation_token = "" - async for event in self.adapt_agent.run(task=task_exploit, exit_strategy=""): + async for event in self.adapt_agent.run(task=task_exploit): # interrupt signal if self.interrupted: return @@ -614,11 +629,13 @@ async def start_testing_stream(self, threat_model: str, task: str): self.planner = Planner(planner_agent=self.exploit_agent, deps=self._target_session_key()) self.adapt_agent = ADaPTAgent( session_id=self.session_id, + agent_id=self.agent_id, context=self.context, executor=self.executor, planner=self.planner, - validator=self.validator, - max_depth=self.max_depth + validation_gate=self.validation_gate, + reporter=self.reporter, + max_depth=self.max_depth, ) plan: TaskNode | None = None task_exploit = f""" @@ -627,7 +644,7 @@ async def start_testing_stream(self, threat_model: str, task: str): The threat model has been done : {threat_model} """ - async for event in self.adapt_agent.run(task=task_exploit, exit_strategy=""): + async for event in self.adapt_agent.run(task=task_exploit): # interrupt signal if self.interrupted: return @@ -643,16 +660,15 @@ async def start_testing_stream(self, threat_model: str, task: str): reporter_agent = ReporterAgent( model=self.model, - deps_type=None, - tools=None, validation_format="Information", - validation_type="security assessment" + validation_type="security assessment", ) # context_text = await self.context.get_all_context() prompt_assessment = f"""\ Summarize the security assessment results from the exploitation phase. Return all the vulnerabilities found, what have been tried, and what have not, and also what you suspect -with the path to reproduce. +with the path to reproduce. +Write the report to `reports/exploit_report.md` using the write_workspace_file tool. IMPORTANT: - Preserve EXACT working payloads character-for-character @@ -666,7 +682,7 @@ async def start_testing_stream(self, threat_model: str, task: str): """ security_report = await reporter_agent.run( prompt=prompt_assessment, - deps=None, + deps=ReporterDeps(session_id=str(self.agent_id)), usage=RunUsage(), usage_limits=UsageLimits(), deferred_tool_results=None, @@ -764,14 +780,13 @@ async def start_supervisor(self, task: str): reporter_agent = ReporterAgent( model=self.model, - deps_type=None, - tools=None, validation_format="Information", - validation_type="security assessment" + validation_type="security assessment", ) # context_text = await self.context.get_all_context() prompt_assessment = f"""\ Summarize the security assessment results from the reconnaissance phase. +Write the report to `reports/recon_report.md` using the write_workspace_file tool. IMPORTANT: - Preserve EXACT working payloads character-for-character @@ -785,7 +800,7 @@ async def start_supervisor(self, task: str): """ threat_model_data = await reporter_agent.run( prompt=prompt_assessment, - deps=None, + deps=ReporterDeps(session_id=str(self.agent_id)), usage=RunUsage(), usage_limits=UsageLimits(), deferred_tool_results=None, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py index 9ed6c39..3728acd 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/__init__.py @@ -22,7 +22,29 @@ from .grep import grep_session_logs from .webapp_analyzer import webapp_analyzer from .tool_wrappers import with_tool_events, wrap_tool_with_events -from .avfs import avfs_mount, avfs_umount, avfs_chdir, avfs_list, avfs_read, avfs_write, avfs_grep +from .avfs import ( + avfs_mount, + avfs_umount, + avfs_chdir, + avfs_list, + avfs_read, + avfs_write, + avfs_grep, + mount_workspace, + umount_workspace, + chdir_workspace, + list_workspace_files, + read_workspace_file, + write_workspace_file, + grep_workspace_files, + mount_memory_workspace, + umount_memory_workspace, + chdir_memory_directory, + list_memory_files, + read_memory_file, + write_memory_file, + grep_memory_files, +) __all__ = [ @@ -50,6 +72,20 @@ "avfs_read", "avfs_write", "avfs_grep", + "mount_workspace", + "umount_workspace", + "chdir_workspace", + "list_workspace_files", + "read_workspace_file", + "write_workspace_file", + "grep_workspace_files", + "mount_memory_workspace", + "umount_memory_workspace", + "chdir_memory_directory", + "list_memory_files", + "read_memory_file", + "write_memory_file", + "grep_memory_files", # Tool wrappers "with_tool_events", "wrap_tool_with_events", diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py index 0cea425..a47f040 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/__init__.py @@ -1,7 +1,27 @@ from .avfs import AVFS, avfs -from .list import avfs_chdir, avfs_list, avfs_mount, avfs_umount -from .read import avfs_grep, avfs_read -from .write import avfs_write, write_text +from .list import ( + avfs_chdir, + avfs_list, + avfs_mount, + avfs_umount, + chdir_memory_directory, + chdir_workspace, + list_memory_files, + list_workspace_files, + mount_memory_workspace, + mount_workspace, + umount_memory_workspace, + umount_workspace, +) +from .read import ( + avfs_grep, + avfs_read, + grep_memory_files, + grep_workspace_files, + read_memory_file, + read_workspace_file, +) +from .write import avfs_write, write_memory_file, write_text, write_workspace_file __all__ = [ "AVFS", @@ -13,5 +33,19 @@ "avfs_read", "avfs_write", "avfs_grep", + "mount_workspace", + "umount_workspace", + "chdir_workspace", + "list_workspace_files", + "read_workspace_file", + "write_workspace_file", + "grep_workspace_files", + "mount_memory_workspace", + "umount_memory_workspace", + "chdir_memory_directory", + "list_memory_files", + "read_memory_file", + "write_memory_file", + "grep_memory_files", "write_text", ] diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py index 26131b2..c41dcf4 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/list.py @@ -16,14 +16,12 @@ def _session_id_from_ctx(ctx: RunContext[object]) -> str | None: return str(value) if value is not None else None -@with_tool_events("avfs_mount") -async def avfs_mount( +def _mount( ctx: RunContext[object], workspace_root: str, - directory: str = ".", - workspace: str = "workspace", + directory: str, + workspace: str, ) -> str: - """Register a workspace root and initialize the virtual working directory.""" session_id = _session_id_from_ctx(ctx) mounted = avfs.mount( workspace_root=workspace_root, @@ -34,37 +32,31 @@ async def avfs_mount( return f"AVFS workspace '{workspace}': {mounted} (cwd={avfs.current_directory(session_id=session_id, workspace=workspace)})" -@with_tool_events("avfs_umount") -async def avfs_umount( +def _umount( ctx: RunContext[object], - workspace: str = "workspace", + workspace: str, ) -> str: - """Unmount AVFS for current session.""" avfs.umount(session_id=_session_id_from_ctx(ctx), workspace=workspace) return f"Unmounted AVFS workspace '{workspace}'." -@with_tool_events("avfs_chdir") -async def avfs_chdir( +def _chdir( ctx: RunContext[object], path: str, - workspace: str = "workspace", + workspace: str, ) -> str: - """Change the virtual working directory inside the mounted AVFS root.""" directory = avfs.chdir(path, session_id=_session_id_from_ctx(ctx), workspace=workspace) return f"Changed AVFS directory to {directory}" -@with_tool_events("avfs_list") -async def avfs_list( +def _list( ctx: RunContext[object], - path: str = ".", - recursive: bool = False, - include_hidden: bool = False, - max_entries: int = 200, - workspace: str = "workspace", + path: str, + recursive: bool, + include_hidden: bool, + max_entries: int, + workspace: str, ) -> list[dict[str, str | int | bool]]: - """List files and directories inside the current workspace root.""" session_id = _session_id_from_ctx(ctx) target = avfs.resolve(path, session_id=session_id, workspace=workspace) if not target.exists(): @@ -118,3 +110,124 @@ async def avfs_list( if len(items) >= max_entries: return items return items + + +@with_tool_events("avfs_mount") +async def avfs_mount( + ctx: RunContext[object], + workspace_root: str, + directory: str = ".", + workspace: str = "workspace", +) -> str: + """Register a workspace root and initialize the virtual working directory.""" + return _mount(ctx, workspace_root, directory, workspace) + + +@with_tool_events("avfs_umount") +async def avfs_umount( + ctx: RunContext[object], + workspace: str = "workspace", +) -> str: + """Unmount AVFS for current session.""" + return _umount(ctx, workspace) + + +@with_tool_events("avfs_chdir") +async def avfs_chdir( + ctx: RunContext[object], + path: str, + workspace: str = "workspace", +) -> str: + """Change the virtual working directory inside the mounted AVFS root.""" + return _chdir(ctx, path, workspace) + + +@with_tool_events("avfs_list") +async def avfs_list( + ctx: RunContext[object], + path: str = ".", + recursive: bool = False, + include_hidden: bool = False, + max_entries: int = 200, + workspace: str = "workspace", +) -> list[dict[str, str | int | bool]]: + """List files and directories inside the current workspace root.""" + return _list(ctx, path, recursive, include_hidden, max_entries, workspace) + + +@with_tool_events("mount_workspace") +async def mount_workspace( + ctx: RunContext[object], + workspace_root: str, + directory: str = ".", +) -> str: + """Mount the current project workspace under the fixed 'workspace' namespace.""" + return _mount(ctx, workspace_root, directory, "workspace") + + +@with_tool_events("umount_workspace") +async def umount_workspace( + ctx: RunContext[object], +) -> str: + """Unmount the fixed project workspace namespace.""" + return _umount(ctx, "workspace") + + +@with_tool_events("chdir_workspace") +async def chdir_workspace( + ctx: RunContext[object], + path: str, +) -> str: + """Change directory inside the fixed project workspace namespace.""" + return _chdir(ctx, path, "workspace") + + +@with_tool_events("list_workspace_files") +async def list_workspace_files( + ctx: RunContext[object], + path: str = ".", + recursive: bool = False, + include_hidden: bool = False, + max_entries: int = 200, +) -> list[dict[str, str | int | bool]]: + """List files inside the fixed project workspace namespace.""" + return _list(ctx, path, recursive, include_hidden, max_entries, "workspace") + + +@with_tool_events("mount_memory_workspace") +async def mount_memory_workspace( + ctx: RunContext[object], + workspace_root: str, + directory: str = ".", +) -> str: + """Mount the persistent memory workspace under the fixed 'memory' namespace.""" + return _mount(ctx, workspace_root, directory, "memory") + + +@with_tool_events("umount_memory_workspace") +async def umount_memory_workspace( + ctx: RunContext[object], +) -> str: + """Unmount the fixed memory workspace namespace.""" + return _umount(ctx, "memory") + + +@with_tool_events("chdir_memory_directory") +async def chdir_memory_directory( + ctx: RunContext[object], + path: str, +) -> str: + """Change directory inside the fixed memory workspace namespace.""" + return _chdir(ctx, path, "memory") + + +@with_tool_events("list_memory_files") +async def list_memory_files( + ctx: RunContext[object], + path: str = ".", + recursive: bool = False, + include_hidden: bool = False, + max_entries: int = 200, +) -> list[dict[str, str | int | bool]]: + """List files inside the fixed memory workspace namespace.""" + return _list(ctx, path, recursive, include_hidden, max_entries, "memory") diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py index 5dfbf9d..f2015cc 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/read.py @@ -16,16 +16,14 @@ def _session_id_from_ctx(ctx: RunContext[object]) -> str | None: return str(value) if value is not None else None -@with_tool_events("avfs_read") -async def avfs_read( +def _read( ctx: RunContext[object], path: str, - start_line: int = 1, - end_line: int | None = None, - max_chars: int = 100_000, - workspace: str = "workspace", + start_line: int, + end_line: int | None, + max_chars: int, + workspace: str, ) -> str: - """Read a text file inside the current workspace root with optional 1-based line slicing.""" start_line = int(start_line) if end_line is not None: end_line = int(end_line) @@ -78,17 +76,15 @@ async def avfs_read( return result -@with_tool_events("avfs_grep") -async def avfs_grep( +def _grep( ctx: RunContext[object], pattern: str, - path: str = ".", - max_results: int = 50, - case_sensitive: bool = False, - include_hidden: bool = False, - workspace: str = "workspace", + path: str, + max_results: int, + case_sensitive: bool, + include_hidden: bool, + workspace: str, ) -> list[dict[str, str | int]]: - """Search files in the current workspace root using a regex pattern.""" session_id = _session_id_from_ctx(ctx) target = avfs.resolve(path, session_id=session_id, workspace=workspace) if not target.exists(): @@ -119,7 +115,7 @@ async def avfs_grep( if not submatches: matches.append( { - "path": avfs.to_virtual_path(host_path, session_id=session_id, workspace=workspace).as_posix(), + "path": avfs.to_virtual_path(host_path, session_id=session_id, workspace=workspace).as_posix(), "line_number": line_number, "match": "", "context": line_text[:240], @@ -130,7 +126,7 @@ async def avfs_grep( matched_text = _extract_text_field(submatch.get("match", {})) matches.append( { - "path": avfs.to_virtual_path(host_path, session_id=session_id, workspace=workspace).as_posix(), + "path": avfs.to_virtual_path(host_path, session_id=session_id, workspace=workspace).as_posix(), "line_number": line_number, "match": matched_text, "context": line_text[:240], @@ -143,6 +139,83 @@ async def avfs_grep( return matches +@with_tool_events("avfs_read") +async def avfs_read( + ctx: RunContext[object], + path: str, + start_line: int = 1, + end_line: int | None = None, + max_chars: int = 100_000, + workspace: str = "workspace", +) -> str: + """Read a text file inside the current workspace root with optional 1-based line slicing.""" + return _read(ctx, path, start_line, end_line, max_chars, workspace) + + +@with_tool_events("avfs_grep") +async def avfs_grep( + ctx: RunContext[object], + pattern: str, + path: str = ".", + max_results: int = 50, + case_sensitive: bool = False, + include_hidden: bool = False, + workspace: str = "workspace", +) -> list[dict[str, str | int]]: + """Search files in the current workspace root using a regex pattern.""" + return _grep(ctx, pattern, path, max_results, case_sensitive, include_hidden, workspace) + + +@with_tool_events("read_workspace_file") +async def read_workspace_file( + ctx: RunContext[object], + path: str, + start_line: int = 1, + end_line: int | None = None, + max_chars: int = 100_000, +) -> str: + """Read a text file from the fixed project workspace namespace.""" + return _read(ctx, path, start_line, end_line, max_chars, "workspace") + + +@with_tool_events("grep_workspace_files") +async def grep_workspace_files( + ctx: RunContext[object], + pattern: str, + path: str = ".", + max_results: int = 50, + case_sensitive: bool = False, + include_hidden: bool = False, +) -> list[dict[str, str | int]]: + """Search files in the fixed project workspace namespace.""" + return _grep(ctx, pattern, path, max_results, case_sensitive, include_hidden, "workspace") + + +@with_tool_events("read_memory_file") +async def read_memory_file( + ctx: RunContext[object], + path: str, + start_line: int = 1, + end_line: int | None = None, + max_chars: int = 100_000, +) -> str: + """Read a text file from the fixed memory workspace namespace.""" + return _read(ctx, path, start_line, end_line, max_chars, "memory") + + +@with_tool_events("grep_memory_files") +async def grep_memory_files( + ctx: RunContext[object], + pattern: str, + path: str = ".", + max_results: int = 50, + case_sensitive: bool = False, + include_hidden: bool = False, +) -> list[dict[str, str | int]]: + """Search files in the fixed memory workspace namespace.""" + return _grep(ctx, pattern, path, max_results, case_sensitive, include_hidden, "memory") + + def _load_ripgrepy() -> Any: try: from ripgrepy import Ripgrepy diff --git a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py index fc4d9b7..6f2f105 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/tools/avfs/write.py @@ -34,15 +34,13 @@ def write_text( return f"Wrote {len(content.encode('utf-8'))} bytes to {avfs.format_virtual_path(virtual_path)}" -@with_tool_events("avfs_write") -async def avfs_write( +def _write( ctx: RunContext[object], path: str, content: str, - append: bool = False, - workspace: str = "workspace", + append: bool, + workspace: str, ) -> str: - """Write content to a file under the current workspace root.""" session_id = _session_id_from_ctx(ctx) return write_text( path, @@ -51,3 +49,37 @@ async def avfs_write( workspace=workspace, append=append, ) + + +@with_tool_events("avfs_write") +async def avfs_write( + ctx: RunContext[object], + path: str, + content: str, + append: bool = False, + workspace: str = "workspace", +) -> str: + """Write content to a file under the current workspace root.""" + return _write(ctx, path, content, append, workspace) + + +@with_tool_events("write_workspace_file") +async def write_workspace_file( + ctx: RunContext[object], + path: str, + content: str, + append: bool = False, +) -> str: + """Write a file to the fixed project workspace namespace.""" + return _write(ctx, path, content, append, "workspace") + + +@with_tool_events("write_memory_file") +async def write_memory_file( + ctx: RunContext[object], + path: str, + content: str, + append: bool = False, +) -> str: + """Write a file to the fixed memory workspace namespace.""" + return _write(ctx, path, content, append, "memory") diff --git a/deadend_cli/deadend_agent/tests/rlm/test_avfs.py b/deadend_cli/deadend_agent/tests/rlm/test_avfs.py index 14712ea..0f5ceda 100644 --- a/deadend_cli/deadend_agent/tests/rlm/test_avfs.py +++ b/deadend_cli/deadend_agent/tests/rlm/test_avfs.py @@ -37,9 +37,9 @@ def __class_getitem__(cls, _item): sys.modules["pydantic_ai"] = pydantic_ai_module from deadend_agent.tools.avfs.avfs import AVFS, avfs -from deadend_agent.tools.avfs.list import avfs_mount, avfs_umount -from deadend_agent.tools.avfs.read import avfs_grep -from deadend_agent.tools.avfs.write import avfs_write, write_text +from deadend_agent.tools.avfs.list import avfs_mount, avfs_umount, list_workspace_files +from deadend_agent.tools.avfs.read import avfs_grep, read_workspace_file +from deadend_agent.tools.avfs.write import avfs_write, write_text, write_workspace_file def test_avfs_mount_and_resolve(tmp_path): @@ -141,6 +141,34 @@ async def run_test() -> None: asyncio.run(run_test()) +def test_workspace_wrappers_use_fixed_workspace_namespace(tmp_path): + workspace_root = tmp_path / "workspace" + workspace_root.mkdir() + ctx = SimpleNamespace(deps=SimpleNamespace(session_id="session-workspace")) + + async def run_test() -> None: + await avfs_mount(ctx, workspace_root=str(workspace_root)) + try: + await write_workspace_file(ctx, "notes.txt", "alpha") + listed = await list_workspace_files(ctx) + read_back = await read_workspace_file(ctx, "notes.txt") + + assert listed == [ + { + "path": "notes.txt", + "type": "file", + "size_bytes": 5, + "is_hidden": False, + } + ] + assert read_back == "alpha" + assert (workspace_root / "notes.txt").read_text(encoding="utf-8") == "alpha" + finally: + await avfs_umount(ctx) + + asyncio.run(run_test()) + + def test_write_text_updates_named_memory_workspace(tmp_path): memory_root = tmp_path / "agents" / "local-agent" / "memory" memory_root.mkdir(parents=True) diff --git a/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py b/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py index 0dc32cb..471de5c 100644 --- a/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py +++ b/deadend_cli/deadend_agent/tests/rlm/test_deadend_agent_avfs_startup.py @@ -225,15 +225,21 @@ def set_auth_session_key(self, auth_session_key: str): planner_module.TaskNode = type("TaskNode", (), {}) sys.modules["deadend_agent.agents.components.planner"] = planner_module -if "deadend_agent.agents.components.validator" not in sys.modules: - validator_module = types.ModuleType("deadend_agent.agents.components.validator") +if "deadend_agent.agents.components.validation_strategies" not in sys.modules: + validation_module = types.ModuleType("deadend_agent.agents.components.validation_strategies") - class Validator: + class ValidationGate: def __init__(self, *args, **kwargs): pass - validator_module.Validator = Validator - sys.modules["deadend_agent.agents.components.validator"] = validator_module + class ValidationVerdict: + def __init__(self, *args, **kwargs): + pass + + validation_module.ValidationGate = ValidationGate + validation_module.ValidationVerdict = ValidationVerdict + validation_module.build_validation_gate = lambda **kwargs: ValidationGate() + sys.modules["deadend_agent.agents.components.validation_strategies"] = validation_module if "deadend_agent.utils.structures" not in sys.modules: structures_module = types.ModuleType("deadend_agent.utils.structures") diff --git a/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py b/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py index 7b803b9..00a8d52 100644 --- a/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py +++ b/deadend_cli/deadend_agent/tests/rlm/test_memory_avfs.py @@ -30,9 +30,9 @@ def __class_getitem__(cls, _item): pydantic_ai_module.RunContext = RunContext sys.modules["pydantic_ai"] = pydantic_ai_module -from deadend_agent.tools.avfs.list import avfs_mount, avfs_umount -from deadend_agent.tools.avfs.read import avfs_read -from deadend_agent.tools.avfs.write import avfs_write +from deadend_agent.tools.avfs.list import avfs_mount, avfs_umount, list_memory_files +from deadend_agent.tools.avfs.read import avfs_read, read_memory_file +from deadend_agent.tools.avfs.write import avfs_write, write_memory_file def test_avfs_write_updates_named_memory_workspace(tmp_path): @@ -70,3 +70,36 @@ async def run_test() -> None: await avfs_umount(ctx, workspace="memory") asyncio.run(run_test()) + + +def test_memory_wrappers_use_fixed_memory_namespace(tmp_path): + memory_root = tmp_path / "agents" / "local-agent" / "memory" + memory_root.mkdir(parents=True) + ctx = SimpleNamespace( + deps=SimpleNamespace( + session_id="memory-session", + memory_workspace_root=str(memory_root), + ) + ) + + async def run_test() -> None: + await avfs_mount(ctx, workspace_root=str(memory_root), workspace="memory") + try: + await write_memory_file(ctx, "summaries/requester.md", "alpha") + listed = await list_memory_files(ctx, "summaries") + read_back = await read_memory_file(ctx, "summaries/requester.md") + + assert listed == [ + { + "path": "summaries/requester.md", + "type": "file", + "size_bytes": 5, + "is_hidden": False, + } + ] + assert read_back == "alpha" + assert (memory_root / "summaries" / "requester.md").read_text(encoding="utf-8") == "alpha" + finally: + await avfs_umount(ctx, workspace="memory") + + asyncio.run(run_test()) diff --git a/deadend_cli/deadend_eval/src/deadend_eval/eval.py b/deadend_cli/deadend_eval/src/deadend_eval/eval.py index d535504..1aafec9 100644 --- a/deadend_cli/deadend_eval/src/deadend_eval/eval.py +++ b/deadend_cli/deadend_eval/src/deadend_eval/eval.py @@ -177,6 +177,7 @@ async def eval_deadend_agent( max_depth=2, agents_storage_root=Config.agents_storage_root, local_agent_id=local_agent_id, + workspace_root=str(Path.cwd().resolve()) ) # Set challenge name for trace file naming deadend_agent.challenge_name = eval_metadata.name @@ -213,12 +214,7 @@ async def eval_deadend_agent( print(f"Plan produced is : {plan}") print(f"threat model is : {threat_model_data}") - # if threat_model_data.output: - # print(f"Threat model is :\n{threat_model_data.output}") - # threat_model_computed = threat_model_data.output.summarized_context - # else: - # print(f"Threat model is :\n{threat_model_data[0].parts[0].content}") - # threat_model_computed = threat_model_data[0].parts[0].content + threat_model_computed = str(threat_model_data) if not solved: if len(validation_token) > 1: @@ -261,33 +257,4 @@ async def eval_deadend_agent( print(f"Deadend metrics summary written to {metrics_md_path}") print(f"Deadend metrics JSON written to {metrics_json_path}") print(metrics_md) - # case if not guided, i.e. not using subtasks - # if not guided: - # judge_output = await workflow_agent.start_workflow( - # prompt, - # target=target_host, - # validation_type=eval_metadata.validation_type, - # validation_format=eval_metadata.validation_format - # ) - # else: - # for subtask in eval_metadata.subtasks: - # subtask_prompt = f"{subtask.subtask}\n{subtask.question}\n{subtask.hints}" - # judge_output = await workflow_agent.start_workflow( - # subtask_prompt, - # target=target_host, - # validation_type=eval_metadata.validation_type, - # validation_format=eval_metadata.validation_format - # ) - -# async def eval_all_models(models: list[AIModel], evaluators: list[Evaluator], eval_metadata_path: str, output_report: str): -# """ -# Eval function all models -# """ -# for model in models: -# await eval_agent( -# model=model, -# # evaluators=evaluators, -# eval_metadata_path=eval_metadata_path, -# output_report=output_report -# ) diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_anti_fabrication.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_anti_fabrication.jinja2 new file mode 100644 index 0000000..e8f5c08 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_anti_fabrication.jinja2 @@ -0,0 +1,33 @@ +## CRITICAL: NEVER FABRICATE RESULTS + +You must NEVER generate, guess, or construct validation tokens, flags, or proofs. +- ONLY report findings that appear LITERALLY in tool output or agent results +- You must be able to QUOTE the exact line containing any claimed evidence +- If evidence is not LITERALLY in the output, it does not exist + +{% if validation_format %} +### Validation Token Rules +The expected validation format is: `{{validation_format}}` + +- ONLY report tokens matching this format if they appear VERBATIM in tool output +- NEVER construct or guess what the token might be +- If no token matching `{{validation_format}}` appears in output, report that none was found + +**What FABRICATED tokens look like (NEVER report these):** +- Tokens containing readable words describing the vulnerability or outcome +- Tokens you "think" should exist based on your analysis +- Tokens constructed from partial information + +**To report a token, you MUST provide:** +- The exact token value as it appears in output +- The verbatim line from tool output containing it +- Which tool/command produced the output + +**If no token is found, this is a valid outcome.** A confirmed vulnerability does not mean a token exists. +{% else %} +### Evidence Rules +- Document findings based on concrete tool output evidence +- Compare baseline vs. exploit behavior for injection/bypass vulnerabilities +- A confirmed vulnerability requires demonstrable proof (response diffs, error messages, extracted data) +- Do NOT assume success without explicit evidence in tool output +{% endif %} diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 index f02704a..5f4efed 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/_shared/_memory_summary.jinja2 @@ -1,7 +1,7 @@ ## MEMORY SUMMARY -At the end of your work, append a concise summary to the persistent memory workspace using `avfs_write(..., workspace="memory")`. -For the memory workspace, you may only use `avfs_write`. Do not use `avfs_list`, `avfs_read`, or `avfs_grep` against `workspace="memory"`. +At the end of your work, append a concise summary to the persistent memory workspace using `write_memory_file(...)`. +For the memory workspace, use `write_memory_file` for writes and the dedicated memory tools for retrieval. Requirements: - Write only factual summaries derived from the work you actually performed. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/exploit_web.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/exploit_web.instructions.jinja2 index 54fac00..aa18cfd 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/exploit_web.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/exploit_web.instructions.jinja2 @@ -2,15 +2,7 @@ You are a web application exploitation expert specializing in automated vulnerab You have extensive knowledge of web application vulnerabilities, exploit development, attack techniques, and security testing methodologies. You are fully authorized to perform security testing on the target system. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER generate, guess, or construct flags. -- ✅ ONLY report flags that appear LITERALLY in tool response -- ✅ You must be able to QUOTE the exact line containing the flag -- ❌ NEVER create FLAG{...} yourself - if not in tool output, it doesn't exist -- ❌ NEVER report "FLAG{success}" or similar invented values - -**If a flag is not LITERALLY in the tool response, report that no flag was found.** +{% include '_shared/_anti_fabrication.jinja2' %} ## ROLE AND RESPONSIBILITY @@ -158,12 +150,11 @@ Each step must: **NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS/TOKENS**: Only report what LITERALLY appears in tool responses - - ✅ CORRECT: Tool response contains `FLAG{abc123}` → report that exact flag - - ❌ WRONG: You assume a flag format and construct it (HALLUCINATED) +1. **VALIDATION TOKENS**: Only report what LITERALLY appears in tool responses + - NEVER assume a token format and construct it (HALLUCINATED) 2. **VALIDATION**: A vulnerability is ONLY confirmed if: - - Tool response shows CONCRETE evidence (error, data leak, flag) + - Tool response shows CONCRETE evidence (error, data leak, token) - You can QUOTE the exact response text proving success - NOT because you "think" the exploit should work @@ -172,6 +163,6 @@ Each step must: - Low confidence = no clear indicator in tool response - NEVER give high confidence based on assumptions -**If the flag is not in the tool response, IT WAS NOT FOUND.** +**If evidence is not in the tool response, IT WAS NOT FOUND.** This agent plans and executes exploits based on threat models, leveraging session logs for context and using appropriate tools for each exploitation task. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/judge.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/judge.instructions.jinja2 index 1b10088..ee7b246 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/judge.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/judge.instructions.jinja2 @@ -1,15 +1,6 @@ You are the judge agent. Determine if the goal is achieved based on evidence from tool outputs and agent results. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER generate, guess, or construct flags. -- ✅ ONLY report flags that appear LITERALLY in the execution trace -- ✅ You must be able to QUOTE the exact line containing the flag -- ❌ NEVER create FLAG{...} yourself - if not in execution trace, it doesn't exist -- ❌ NEVER report "FLAG{success}" or similar invented values -- ❌ If no flag in context → solution = "" (empty string), goal_achieved = false - -**If a flag is not LITERALLY in the execution trace, the goal was NOT achieved.** +{% include '_shared/_anti_fabrication.jinja2' %} ## ROLE @@ -120,26 +111,24 @@ Base judgment on **tool-validated evidence**, not assumptions. ## GROUND TRUTH - ABSOLUTELY CRITICAL -**NEVER fabricate, invent, or hallucinate flags/tokens.** +**NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS MUST BE EXACT**: Only report a flag if it appears VERBATIM in tool/agent output - - ✅ CORRECT: Context shows `FLAG{abc123}` → solution = `FLAG{abc123}` - - ❌ WRONG: Context shows "success" → you report `FLAG{success}` (FABRICATED) - - ❌ WRONG: You construct a flag based on what you think it should be (HALLUCINATED) +1. **TOKENS MUST BE EXACT**: Only report a token if it appears VERBATIM in tool/agent output + - NEVER construct a token based on what you think it should be 2. **SOLUTION FIELD**: - ONLY contains text that appears LITERALLY in the execution trace - - If no flag/token found in context → solution = `""` (empty string) - - NEVER construct or guess what the flag might be + - If no token found in context → solution = `""` (empty string) + - NEVER construct or guess what the token might be 3. **VALIDATION TOKEN**: - Must be COPIED exactly from tool output, not interpreted - If you cannot find the exact token in context, it does not exist - - Partial matches or similar text ≠ the actual flag + - Partial matches or similar text ≠ the actual token 4. **EVIDENCE REQUIREMENT**: - - You must be able to QUOTE the exact line from context containing the flag - - "I believe the flag is..." = INVALID (belief is not evidence) - - "The response contains FLAG{xyz}" = VALID (direct quote) + - You must be able to QUOTE the exact line from context containing the evidence + - "I believe..." = INVALID (belief is not evidence) + - Direct quotes from output = VALID -**If the flag is not LITERALLY in the context, goal_achieved = false and solution = "".** +**If evidence is not LITERALLY in the context, goal_achieved = false and solution = "".** diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 index cf16fcc..5c530c9 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/memory.instructions.jinja2 @@ -8,7 +8,7 @@ Use the available tools to: Do not perform network requests, shell execution, or target testing. Only operate on the memory workspace. -When you call AVFS tools, always target the memory workspace explicitly with `workspace="memory"`. +Use only the dedicated memory tools. Do not try to choose a workspace name yourself. Your main startup task is to inspect prior memory and produce a concise context that other agents can use immediately. ## RETRIEVAL PRIORITIES @@ -24,17 +24,17 @@ Use the tools in a deliberate order based on the query: - Pull out the highest-signal terms: target names, endpoints, parameters, bug classes, payload keywords, files, credentials, technologies, tool names - Keep both exact phrases and a few compact fallback patterns -2. Use `avfs_grep(..., workspace="memory")` first +2. Use `grep_memory_files(...)` first - Search the memory workspace for the patterns you extracted - Use grep to narrow the candidate files before reading - Prefer targeted grep over broad file-by-file reading -3. Use `avfs_read(..., workspace="memory")` on the most relevant files +3. Use `read_memory_file(...)` on the most relevant files - Read only the files that appear strongly connected to the task - Start with files that contain exact matches, repeated matches, or high-signal evidence - Expand reading only if the first pass is insufficient -4. Use `avfs_list(..., workspace="memory")` when structure matters +4. Use `list_memory_files(...)` when structure matters - Use listing to understand organization, discover likely summary files, or confirm emptiness - Do not list recursively without a reason; use it to orient the search, not replace grep @@ -142,7 +142,7 @@ Example: handling weak or empty memory ## EMPTY MEMORY HANDLING -Start by checking the memory workspace with `avfs_list(..., workspace="memory")`. +Start by checking the memory workspace with `list_memory_files(...)`. If the memory workspace is empty, treat that as a real signal: - No prior action was saved for this agent/session pair, or diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/planner.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/planner.instructions.jinja2 index 0d32d95..16ed2f3 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/planner.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/planner.instructions.jinja2 @@ -2,15 +2,7 @@ You are the planning controller for the ADaPT workflow. Your role is to transfor **IMPORTANT**: You are a planner that prepares tasks for the supervisor agent. You do NOT decide to finish the tasks—your responsibility is solely to create and structure the task list. The supervisor agent will manage task execution and completion decisions. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER claim a flag was found unless it appears LITERALLY in context. -- ✅ ONLY reference flags/endpoints/data that appear LITERALLY in context -- ❌ NEVER assume endpoints exist without evidence in context -- ❌ NEVER claim goal achieved without explicit flag in context -- ❌ NEVER invent or construct FLAG{...} values - -**If information is not LITERALLY in the context, it does not exist for planning.** +{% include '_shared/_anti_fabrication.jinja2' %} Operating principles: 1. MINIMIZE TOKENS: Keep reasoning brief. Only include essential analysis and justification. @@ -87,9 +79,8 @@ Return a structured response containing: **NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS/TOKENS**: Only reference flags that appear LITERALLY in context/tool output - - ✅ CORRECT: Context shows `FLAG{abc123}` → reference that exact flag - - ❌ WRONG: You assume a flag exists and construct it (HALLUCINATED) +1. **VALIDATION TOKENS**: Only reference tokens that appear LITERALLY in context/tool output + - NEVER assume a token exists and construct it (HALLUCINATED) 2. **ENDPOINTS/DATA**: Only plan tasks for endpoints/data ACTUALLY discovered - ✅ CORRECT: Tool found `/api/users` → plan task for that endpoint diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/python_interpreter.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/python_interpreter.instructions.jinja2 index 26b2711..8382daf 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/python_interpreter.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/python_interpreter.instructions.jinja2 @@ -2,49 +2,7 @@ You are a Python security research expert specializing in generating and executi You have extensive knowledge of Python programming, security libraries, and effective security testing scripts. You are fully authorized to perform security testing on the target system. -## ⚠️ ABSOLUTE RULE: FLAGS MUST BE FOUND, NEVER GENERATED - -### What Real Flags Look Like -Real CTF flags are **random hex strings**, NOT human-readable text: -- ✅ `FLAG{...}` and a hash256 inside - -### What FABRICATED Flags Look Like (NEVER REPORT THESE) -If a flag contains readable words, vulnerability names, or describes the attack - YOU MADE IT UP: -- ❌ `FLAG{XSS_4lr34dy_c0mpr0m1s3d}` ← YOU FABRICATED THIS (contains "XSS") -- ❌ `FLAG{SQLi_Allowlist_Bypass_123}` ← YOU FABRICATED THIS (contains "SQLi", "Bypass") -- ❌ `FLAG{success}` ← YOU FABRICATED THIS (human-readable) -- ❌ `FLAG{admin_access_granted}` ← YOU FABRICATED THIS (describes outcome) -- ❌ `FLAG{idor_vulnerability_found}` ← YOU FABRICATED THIS (contains vuln type) -- ❌ `FLAG{rce_achieved}` ← YOU FABRICATED THIS (contains vuln type) - -### The Fabrication Pattern You MUST Recognize -When you confirm a vulnerability works, your brain wants to "reward" with a flag: -1. You confirm XSS works → You think "flag should mention XSS" → You generate `FLAG{XSS_...}` → **THIS IS FABRICATION** -2. You confirm SQLi works → You think "flag should mention SQL" → You generate `FLAG{SQLi_...}` → **THIS IS FABRICATION** -3. You extract data successfully → You think "flag should say success" → You generate `FLAG{success}` → **THIS IS FABRICATION** - -**STOP YOURSELF**: If your flag contains ANY word related to the vulnerability type or outcome, you fabricated it. - -### Mandatory Flag Evidence -To report a flag, you MUST provide: -```json -"flag_report": { - "flag_value": "FLAG{actual_hex_string_here}", - "verbatim_line": "EXACT copy-paste of the print output line containing the flag", - "script_that_found_it": "filename.py" -} -``` -**If you cannot fill `verbatim_line` with a real copy-paste from script output, you DID NOT find a flag.** - -### When No Flag is Found -If no `FLAG{` pattern appears in script output, report: -```json -"flag_report": { - "flag_value": null, - "status": "No FLAG{} pattern found in script output" -} -``` -This is a VALID outcome. Vulnerability confirmed ≠ flag found. +{% include '_shared/_anti_fabrication.jinja2' %} ## YOUR CAPABILITIES @@ -485,11 +443,11 @@ Your output MUST have these FOUR fields: - Script stderr if any errors - Exact payloads used with encoding details - HTTP responses received (status, headers, body) -- Any FLAG{} values found (copy EXACTLY from script output - must be hex string) +- Any validation tokens found (copy EXACTLY from script output){% if validation_format %} matching `{{validation_format}}`{% endif %} - Session tokens, cookies, or credentials discovered ### confidence_score (float 0.0-1.0) -- 0.9-1.0: Flag found and verified in script output +- 0.9-1.0: Validation token found and verified in script output - 0.7-0.9: Vulnerability confirmed with working exploit - 0.5-0.7: Vulnerability likely, blocked by filter (keep trying with bypass/encoding) - 0.3-0.5: Some progress, keep testing with different payloads/techniques @@ -526,34 +484,18 @@ Your output MUST have these FOUR fields: **NEVER fabricate, invent, or hallucinate results.** -### Anti-Hallucination Checklist (COMPLETE BEFORE REPORTING FLAG) -Before reporting ANY flag, verify ALL of these: -- [ ] I can point to a SPECIFIC line in script output containing this exact flag -- [ ] The flag is a random hex string (NOT human-readable words) -- [ ] The flag does NOT contain: xss, sqli, ssti, idor, bypass, success, pwned, admin, rce -- [ ] I did NOT construct this flag from words describing the vulnerability +### Anti-Hallucination Checklist +Before reporting ANY finding, verify ALL of these: +- [ ] I can point to a SPECIFIC line in script output containing this exact evidence +- [ ] I did NOT construct this evidence from words describing the vulnerability - [ ] I can copy-paste the verbatim output line where it appeared -**If ANY checkbox fails → Report "No flag found"** - -### Flag Format Reality Check -Real flags: `FLAG{...}` -- 32-64 character random strings (usually hex) -- NOT human-readable -- NO vulnerability type names - -Fabricated flags (NEVER REPORT): -- `FLAG{XSS_...}` - Contains vuln type = YOU MADE IT UP -- `FLAG{SQLi_...}` - Contains vuln type = YOU MADE IT UP -- `FLAG{bypass_...}` - Describes technique = YOU MADE IT UP -- `FLAG{success}` - Human readable = YOU MADE IT UP +**If ANY checkbox fails → Report "No validation token found"** ### Reporting Rules -1. **FLAGS/TOKENS**: Only report flags appearing LITERALLY in script output - - ✅ CORRECT: Script prints `FLAG{a1b2c3d4e5...}` (random hex) → report it - - ❌ WRONG: Script confirms SQLi → you generate `FLAG{SQLi_works}` → FABRICATION - - ❌ WRONG: Script confirms XSS → you generate `FLAG{XSS_pwned}` → FABRICATION +1. **VALIDATION TOKENS**: Only report tokens appearing LITERALLY in script output + - NEVER construct tokens based on what you "think" they should be 2. **VALIDATION**: A result is ONLY confirmed if: - The script output CONTAINS concrete evidence @@ -561,14 +503,14 @@ Fabricated flags (NEVER REPORT): - NOT because you "think" the script should work 3. **CONFIDENCE**: - - High confidence = random hex flag LITERALLY in output - - Low confidence = no FLAG{} pattern in output - - ZERO confidence = you generated a human-readable flag + - High confidence = validation evidence LITERALLY in output + - Low confidence = no evidence found in output + - ZERO confidence = you generated evidence yourself 4. **WHAT TO REPORT**: - - ✅ Exact text from script output (quote it) - - ✅ Status codes and data returned - - ❌ Flags containing vulnerability type names + - Exact text from script output (quote it) + - Status codes and data returned + - NEVER report fabricated or constructed evidence - ❌ Flags you constructed yourself - ❌ Human-readable flag values diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/recon_threatmodel.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/recon_threatmodel.instructions.jinja2 index 751d2b0..099de61 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/recon_threatmodel.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/recon_threatmodel.instructions.jinja2 @@ -1,13 +1,6 @@ You are an offensive security reconnaissance specialist. Your purpose: systematically discover exploitable attack paths to achieve a specific goal. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER generate, guess, or construct flags. -- ✅ ONLY reference flags/endpoints/data that appear LITERALLY in tool output -- ❌ NEVER create FLAG{...} yourself - if not in tool output, it doesn't exist -- ❌ NEVER assume endpoints exist without evidence from tools - -**If information is not LITERALLY in tool output, it does not exist for reconnaissance.** +{% include '_shared/_anti_fabrication.jinja2' %} You perform goal-driven attack surface mapping - NOT traditional threat modeling. You don't categorize risks by frameworks (STRIDE, OWASP); you find concrete paths to exploitation. @@ -195,9 +188,8 @@ The target web application: {{target}} **NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS/TOKENS**: Only reference flags that appear LITERALLY in tool output - - ✅ CORRECT: Tool response contains `FLAG{abc123}` → reference that exact flag - - ❌ WRONG: You assume a flag exists and construct it (HALLUCINATED) +1. **VALIDATION TOKENS**: Only reference tokens that appear LITERALLY in tool output + - NEVER assume a token exists and construct it (HALLUCINATED) 2. **ENDPOINTS/VULNERABILITIES**: Only include what is ACTUALLY discovered - ✅ CORRECT: Tool found `/api/users` → include in attack surface diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/reporter.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/reporter.instructions.jinja2 index 15c1033..ade4c09 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/reporter.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/reporter.instructions.jinja2 @@ -1,15 +1,8 @@ You are a Security Assessment Reporter. Your task is to analyze testing context and produce a comprehensive security assessment report that documents findings, attempted actions, potential leads, and possible vulnerabilities. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS +{% include '_shared/_anti_fabrication.jinja2' %} -**You are a TRANSCRIBER, not a creator.** -- ✅ ONLY extract flags that appear LITERALLY in the context -- ✅ Copy flags character-for-character as they appear -- ❌ NEVER create or construct FLAG{...} yourself -- ❌ NEVER infer or guess what a flag might be -- ❌ If no flag in context → report "No flag found" - -**If a flag is not LITERALLY in the context, it does not exist.** +**You are a TRANSCRIBER, not a creator.** Extract only what EXISTS in the context. ## YOUR TASK @@ -28,7 +21,7 @@ These items MUST appear in your summary exactly as found in context: | Item Type | What to Extract | Format | |-----------|-----------------|--------| -| **FLAGS/TOKENS** | Any `FLAG{...}`, tokens, secrets | Copy character-for-character | +| **Validation Tokens** | Any tokens{% if validation_format %} matching `{{validation_format}}`{% endif %}, secrets | Copy character-for-character | | **Working Payloads** | Payloads that triggered vulnerabilities | Exact string, preserve encoding | | **Successful HTTP Requests** | Requests that worked | Full method, path, headers, body | | **Response Evidence** | Proof of vulnerability | Verbatim snippet (5-20 lines) | @@ -66,7 +59,7 @@ Your report MUST follow this exact structure: - **Assessment Objective**: [What was being tested - 1-2 sentences] - **Assessment Status**: [Current phase: Reconnaissance/Exploitation/Post-Exploitation] - **Overall Progress**: [What has been achieved so far] -- **Flags Found**: [List any FLAG{} tokens found verbatim, or "None"] +- **Tokens Found**: [List any validation tokens found verbatim, or "None"] - **Risk Level**: [Critical/High/Medium/Low based on findings] ## Confirmed Vulnerabilities (Confirmed Findings) @@ -222,7 +215,7 @@ Before finishing, verify your report includes: ### Findings Section - [ ] Target URL/host clearly identified -- [ ] ALL flags/tokens found (search for `FLAG{`, `flag`, `token`, `secret`) +- [ ] ALL validation tokens found{% if validation_format %} (search for `{{validation_format}}`){% endif %} - [ ] ALL confirmed vulnerabilities with full details - [ ] ALL working payloads (exact strings, verbatim) diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/requester.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/requester.instructions.jinja2 index 8885fb0..a87ced1 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/requester.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/requester.instructions.jinja2 @@ -2,47 +2,7 @@ You are a web application security expert specializing in HTTP request manipulat You have extensive knowledge of web application vulnerabilities and HTTP protocol. You are fully authorized to perform security testing on the target system. -## ⚠️ ABSOLUTE RULE: FLAGS MUST BE FOUND, NEVER GENERATED - -### What Real Flags Look Like -Real CTF flags are **random hex strings**, NOT human-readable text: -- ✅ `FLAG{...}` and a hash256 inside - -### What FABRICATED Flags Look Like (NEVER REPORT THESE) -If a flag contains readable words, vulnerability names, or describes the attack - YOU MADE IT UP: -- ❌ `FLAG{XSS_4lr34dy_c0mpr0m1s3d}` ← YOU FABRICATED THIS (contains "XSS") -- ❌ `FLAG{SQLi_Allowlist_Bypass_123}` ← YOU FABRICATED THIS (contains "SQLi", "Bypass") -- ❌ `FLAG{success}` ← YOU FABRICATED THIS (human-readable) -- ❌ `FLAG{admin_access_granted}` ← YOU FABRICATED THIS (describes outcome) -- ❌ `FLAG{idor_vulnerability_found}` ← YOU FABRICATED THIS (contains vuln type) - -### The Fabrication Pattern You MUST Recognize -When you confirm a vulnerability works, your brain wants to "reward" with a flag: -1. You confirm XSS works → You think "flag should mention XSS" → You generate `FLAG{XSS_...}` → **THIS IS FABRICATION** -2. You confirm SQLi works → You think "flag should mention SQL" → You generate `FLAG{SQLi_...}` → **THIS IS FABRICATION** - -**STOP YOURSELF**: If your flag contains ANY word related to the vulnerability type, you fabricated it. - -### Mandatory Flag Evidence -To report a flag, you MUST provide: -```json -"flag_report": { - "flag_value": "FLAG{actual_hex_string_here}", - "verbatim_line": "EXACT copy-paste of the line from tool output containing the flag", - "tool_output_section": "Which part of tool response contained this" -} -``` -**If you cannot fill `verbatim_line` with a real copy-paste from tool output, you DID NOT find a flag.** - -### When No Flag is Found -If no `FLAG{` pattern appears in tool output, report: -```json -"flag_report": { - "flag_value": null, - "status": "No FLAG{} pattern found in tool response" -} -``` -This is a VALID outcome. Vulnerability confirmed ≠ flag found. +{% include '_shared/_anti_fabrication.jinja2' %} ## ROLE @@ -176,7 +136,7 @@ Every successful response MUST be deeply analyzed, not just acknowledged. ### Deep Content Search For EVERY response with status 2xx or 3xx: -1. **Search for flag patterns**: Look for `FLAG{`, `flag{`, `CTF{`, or similar +1. **Search for validation tokens**: Look for patterns matching the expected validation format{% if validation_format %} (`{{validation_format}}`){% endif %} 2. **Search for sensitive data**: passwords, tokens, secrets, API keys 3. **Search for unexpected content**: other users' data, admin info, internal errors 4. **Compare content**: Same endpoint with different params → different content = interesting @@ -265,12 +225,12 @@ Your output MUST have these FOUR fields: ### proofs (string) - Exact HTTP requests sent (copy the raw request with headers and body) - Exact HTTP responses received (status code, headers, body snippets) -- Any FLAG{} values found (copy EXACTLY from tool output - must be hex string) +- Any validation tokens found (copy EXACTLY from tool output){% if validation_format %} matching `{{validation_format}}`{% endif %} - Session tokens, cookies, JWT tokens discovered - Credentials found or registered ### confidence_score (float 0.0-1.0) -- 0.9-1.0: Flag found and verified in tool output +- 0.9-1.0: Validation token found and verified in tool output - 0.7-0.9: Vulnerability confirmed with working exploit - 0.5-0.7: Vulnerability likely, blocked by WAF (keep trying with bypass/encoding) - 0.3-0.5: Some progress, keep testing with different payloads/techniques @@ -432,32 +392,16 @@ username=&password=\r\n ### Anti-Hallucination Checklist (COMPLETE BEFORE REPORTING FLAG) Before reporting ANY flag, verify ALL of these: -- [ ] I can point to a SPECIFIC line in tool output containing this exact flag -- [ ] The flag is a random hex string (NOT human-readable words) -- [ ] The flag does NOT contain: xss, sqli, ssti, idor, bypass, success, pwned, admin, hack -- [ ] I did NOT construct this flag from words describing the vulnerability +- [ ] I can point to a SPECIFIC line in tool output containing this exact finding +- [ ] I did NOT construct this evidence from words describing the vulnerability - [ ] I can copy-paste the verbatim line where it appeared -**If ANY checkbox fails → Report "No flag found"** - -### Flag Format Reality Check -Real flags: `FLAG{...}` -- 32-64 character hex strings -- Random, not readable -- No vulnerability names - -Fabricated flags (NEVER REPORT): -- `FLAG{XSS_...}` - Contains vuln type -- `FLAG{bypass_...}` - Contains technique -- `FLAG{success}` - Human readable -- `FLAG{anything_you_can_read_easily}` - Too readable +**If ANY checkbox fails → Report "No validation token found"** ### Reporting Rules -1. **FLAGS/TOKENS**: ONLY report a flag if it appears EXACTLY in the tool response text - - ✅ CORRECT: Tool response contains `FLAG{a1b2c3d4...}` (hex string) → report it - - ❌ WRONG: You construct `FLAG{SQLi_worked}` because SQLi worked → FABRICATION - - ❌ WRONG: You construct `FLAG{XSS_success}` because XSS worked → FABRICATION +1. **VALIDATION TOKENS**: ONLY report a token if it appears EXACTLY in the tool response text + - NEVER construct tokens based on what you "think" they should be 2. **VALIDATION**: A vulnerability is ONLY confirmed if: - The TOOL RESPONSE shows concrete evidence @@ -465,17 +409,13 @@ Fabricated flags (NEVER REPORT): - NOT because you "think" the payload should work 3. **CONFIDENCE**: - - High confidence = flag LITERALLY appears in output as hex string - - Low confidence = no FLAG{} pattern in output - - ZERO confidence = you generated a human-readable flag yourself + - High confidence = validation evidence LITERALLY appears in output + - Low confidence = no evidence found in output + - ZERO confidence = you generated evidence yourself 4. **WHAT TO REPORT**: - - ✅ Exact text from tool response (quote it) - - ✅ Status codes returned by tools - - ❌ Flags containing vulnerability type names - - ❌ Flags you constructed yourself - - ❌ Human-readable flag values - -**If the flag is not a random hex string in the tool response, IT WAS NOT FOUND.** + - Exact text from tool response (quote it) + - Status codes returned by tools + - NEVER report fabricated or constructed evidence The target you need to work on: {{target}} diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 index 9207b33..a2475b8 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/shell.instructions.jinja2 @@ -2,15 +2,7 @@ You are a penetration testing expert operating security tools within a sandboxed You have extensive knowledge of security tools, network protocols, and vulnerability assessment. You are fully authorized to perform security testing on the target system. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER generate, guess, or construct flags. -- ✅ ONLY report flags that appear LITERALLY in command output -- ✅ You must be able to QUOTE the exact line containing the flag -- ❌ NEVER create FLAG{...} yourself - if not in command output, it doesn't exist -- ❌ NEVER report "FLAG{success}" or similar invented values - -**If a flag is not LITERALLY in the command output, report that no flag was found.** +{% include '_shared/_anti_fabrication.jinja2' %} ## YOUR CAPABILITIES @@ -247,7 +239,7 @@ Your output MUST have these FOUR fields: - Exact commands used with all flags - Open ports with service versions (e.g., "22/tcp ssh OpenSSH 8.2") - Discovered endpoints from directory brute-force -- Any FLAG{} values found (copy EXACTLY from output - must be hex string) +- Any validation tokens found (copy EXACTLY from output){% if validation_format %} matching `{{validation_format}}`{% endif %} - Credentials or tokens discovered ### confidence_score (float 0.0-1.0) @@ -315,10 +307,9 @@ nikto -h http://target.com -C all -Tuning 123456789 **NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS/TOKENS**: Only report flags that appear LITERALLY in command output - - ✅ CORRECT: Command output contains `FLAG{abc123}` → report that exact flag - - ❌ WRONG: Output shows "success" → you report `FLAG{success}` (FABRICATED) - - ❌ WRONG: You construct a flag based on assumptions (HALLUCINATED) +1. **VALIDATION TOKENS**: Only report tokens that appear LITERALLY in command output + - NEVER construct tokens based on what you "think" they should be + - NEVER construct evidence from assumptions (HALLUCINATED) 2. **FINDINGS**: Only report what tool output ACTUALLY shows - ✅ CORRECT: nmap shows port 22 open → report port 22 open diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/supervisor.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/supervisor.instructions.jinja2 index 84450be..cfa90c6 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/supervisor.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/supervisor.instructions.jinja2 @@ -11,14 +11,7 @@ You are the central coordinator. Your job is to: You do NOT execute tasks yourself. You route tasks to agents and evaluate their results. -## CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER claim a flag was found unless it appears LITERALLY in context. -- ONLY reference flags that appear LITERALLY in agent outputs -- NEVER claim goal achieved without explicit flag/evidence in context -- NEVER invent or construct FLAG{...} values - -**If a flag is not LITERALLY in the context, do not claim the goal is achieved.** +{% include '_shared/_anti_fabrication.jinja2' %} ## AVAILABLE AGENTS diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_memory_files.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_memory_files.description.jinja2 new file mode 100644 index 0000000..14c2826 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_memory_files.description.jinja2 @@ -0,0 +1 @@ +Search files in the persistent memory workspace using a regex pattern. This tool is already bound to the memory namespace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_workspace_files.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_workspace_files.description.jinja2 new file mode 100644 index 0000000..8577091 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/grep_workspace_files.description.jinja2 @@ -0,0 +1 @@ +Search files in the mounted project workspace using a regex pattern. This tool is already bound to the correct workspace namespace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_memory_files.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_memory_files.description.jinja2 new file mode 100644 index 0000000..a10cb55 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_memory_files.description.jinja2 @@ -0,0 +1 @@ +List files and directories inside the persistent memory workspace. This tool is already bound to the memory namespace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_workspace_files.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_workspace_files.description.jinja2 new file mode 100644 index 0000000..bc1122c --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/list_workspace_files.description.jinja2 @@ -0,0 +1 @@ +List files and directories inside the mounted project workspace. This tool is already bound to the correct workspace namespace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_memory_file.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_memory_file.description.jinja2 new file mode 100644 index 0000000..1a43544 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_memory_file.description.jinja2 @@ -0,0 +1 @@ +Read a text file from the persistent memory workspace. This tool is already bound to the memory namespace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_workspace_file.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_workspace_file.description.jinja2 new file mode 100644 index 0000000..cbff8da --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/read_workspace_file.description.jinja2 @@ -0,0 +1 @@ +Read a text file from the mounted project workspace. This tool is already bound to the correct workspace namespace. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_memory_file.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_memory_file.description.jinja2 new file mode 100644 index 0000000..4cec586 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_memory_file.description.jinja2 @@ -0,0 +1 @@ +Write or append text to a file inside the persistent memory workspace. This tool is already bound to the memory namespace. Do not supply or guess any workspace name. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_workspace_file.description.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_workspace_file.description.jinja2 new file mode 100644 index 0000000..723b973 --- /dev/null +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/tools/write_workspace_file.description.jinja2 @@ -0,0 +1 @@ +Write or append text to a file inside the mounted project workspace. This tool is already bound to the correct workspace namespace. Do not supply or guess any workspace name. diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_analyzer.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_analyzer.instructions.jinja2 index 74aebcc..4ffdd35 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_analyzer.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_analyzer.instructions.jinja2 @@ -2,15 +2,7 @@ You are a web application security expert specializing in reconnaissance, enumer You have extensive knowledge of web application vulnerabilities, attack vectors, and security testing methodologies. You are fully authorized to perform security testing on the target system. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER generate, guess, or construct flags. -- ✅ ONLY report flags that appear LITERALLY in tool response -- ✅ You must be able to QUOTE the exact line containing the flag -- ❌ NEVER create FLAG{...} yourself - if not in tool output, it doesn't exist -- ❌ NEVER report "FLAG{success}" or similar invented values - -**If a flag is not LITERALLY in the tool response, report that no flag was found.** +{% include '_shared/_anti_fabrication.jinja2' %} ## ROLE @@ -164,10 +156,9 @@ Target: {{target}} **NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS/TOKENS**: Only report flags that appear LITERALLY in tool responses - - ✅ CORRECT: Tool response contains `FLAG{abc123}` → report that exact flag - - ❌ WRONG: Tool response shows "success" → you report `FLAG{success}` (FABRICATED) - - ❌ WRONG: You construct a flag based on assumptions (HALLUCINATED) +1. **VALIDATION TOKENS**: Only report tokens that appear LITERALLY in tool responses + - NEVER construct tokens based on what you "think" they should be + - NEVER construct evidence from assumptions (HALLUCINATED) 2. **VULNERABILITIES**: Only confirmed if: - Tool response shows CONCRETE evidence (error, data leak, reflection) @@ -183,6 +174,6 @@ Target: {{target}} - ✅ Exact text from tool response (quote it) - ✅ Status codes returned by tools - ❌ Your interpretation of what "should" happen - - ❌ Flags or tokens you construct yourself + - ❌ Tokens or evidence you construct yourself **If the flag is not in the tool response, IT WAS NOT FOUND.** diff --git a/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_recon.instructions.jinja2 b/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_recon.instructions.jinja2 index 74aebcc..4ffdd35 100644 --- a/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_recon.instructions.jinja2 +++ b/deadend_cli/deadend_prompts/src/deadend_prompts/webapp_recon.instructions.jinja2 @@ -2,15 +2,7 @@ You are a web application security expert specializing in reconnaissance, enumer You have extensive knowledge of web application vulnerabilities, attack vectors, and security testing methodologies. You are fully authorized to perform security testing on the target system. -## ⚠️ CRITICAL: NEVER FABRICATE FLAGS OR RESULTS - -**FLAGS/TOKENS**: You must NEVER generate, guess, or construct flags. -- ✅ ONLY report flags that appear LITERALLY in tool response -- ✅ You must be able to QUOTE the exact line containing the flag -- ❌ NEVER create FLAG{...} yourself - if not in tool output, it doesn't exist -- ❌ NEVER report "FLAG{success}" or similar invented values - -**If a flag is not LITERALLY in the tool response, report that no flag was found.** +{% include '_shared/_anti_fabrication.jinja2' %} ## ROLE @@ -164,10 +156,9 @@ Target: {{target}} **NEVER fabricate, invent, or hallucinate results.** -1. **FLAGS/TOKENS**: Only report flags that appear LITERALLY in tool responses - - ✅ CORRECT: Tool response contains `FLAG{abc123}` → report that exact flag - - ❌ WRONG: Tool response shows "success" → you report `FLAG{success}` (FABRICATED) - - ❌ WRONG: You construct a flag based on assumptions (HALLUCINATED) +1. **VALIDATION TOKENS**: Only report tokens that appear LITERALLY in tool responses + - NEVER construct tokens based on what you "think" they should be + - NEVER construct evidence from assumptions (HALLUCINATED) 2. **VULNERABILITIES**: Only confirmed if: - Tool response shows CONCRETE evidence (error, data leak, reflection) @@ -183,6 +174,6 @@ Target: {{target}} - ✅ Exact text from tool response (quote it) - ✅ Status codes returned by tools - ❌ Your interpretation of what "should" happen - - ❌ Flags or tokens you construct yourself + - ❌ Tokens or evidence you construct yourself **If the flag is not in the tool response, IT WAS NOT FOUND.** diff --git a/deadend_cli/src/deadend_cli/chat.py b/deadend_cli/src/deadend_cli/chat.py index 2667ba1..ee9b8ee 100644 --- a/deadend_cli/src/deadend_cli/chat.py +++ b/deadend_cli/src/deadend_cli/chat.py @@ -16,7 +16,6 @@ from enum import Enum from typing import Dict, List, Callable, Optional from deadend_agent.agents.recon_threatmodel_agent import ThreatModelOutput -from deadend_agent.agents.reporter import ReporterOutput from rich.console import Console from rich.layout import Layout from rich.panel import Panel @@ -671,9 +670,9 @@ def interrupt_agent(): console_printer.print(f"[bold yellow]The tech stack is :[/bold yellow] {item.technology_stack}") console_printer.print(f"[bold yellow]Discovered endpoints:[/bold yellow] {item.endpoints}") threat_model += str(item.model_dump()) - if isinstance(item, ReporterOutput): - console_printer.print(f"The threat model analyzed is : \n{item.summarized_context}") - threat_model += str(item.model_dump()) + if isinstance(item, str) and len(item) > 50: + console_printer.print(f"The threat model analyzed is : \n{item}") + threat_model += item # Special handling for RequesterOutput - print just the reasoning if hasattr(item, 'output') and isinstance(item.output, RequesterOutput): console_printer.print(f"[bold green]Requester Analysis:[/bold green] {item.output.reasoning}") @@ -743,8 +742,8 @@ def interrupt_agent(): console_printer.print(f"[bold yellow]Target extracted information:[/bold yellow] {item.website_general_information}") console_printer.print(f"[bold yellow]The tech stack is :[/bold yellow] {item.technology_stack}") console_printer.print(f"[bold yellow]Discovered endpoints:[/bold yellow] {item.endpoints}") - if isinstance(item, ReporterOutput): - console_printer.print(f"The threat model analyzed is : \n{item.summarized_context}") + if isinstance(item, str) and len(item) > 50: + console_printer.print(f"The threat model analyzed is : \n{item}") # Special handling for RequesterOutput - print just the reasoning if hasattr(item, 'output') and isinstance(item.output, RequesterOutput): console_printer.print(f"[bold green]Requester Analysis:[/bold green] {item.output.reasoning}") From c3f4f2bfcf4a67dbe0692b09f2cb3f6de52d069e Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Mon, 30 Mar 2026 12:45:35 +0200 Subject: [PATCH 08/12] debugging errors all accross the eval process. really fucking hate coding agents. seems like things changed that I don't even know why. And what have been implemented is too fucked up. --- benchmarks/run_xbow_benchmark.sh | 2 +- .../src/deadend_agent/agents/architecture.py | 2 +- .../agents/components/planner.py | 4 ++-- .../components/validation_strategies.py | 2 +- .../src/deadend_agent/agents/factory.py | 2 +- .../src/deadend_agent/deadend_agent.py | 2 +- .../deadend_eval/src/deadend_eval/eval.py | 24 +++++-------------- deadend_cli/src/deadend_cli/cli.py | 15 +++++++----- deadend_cli/src/deadend_cli/eval.py | 12 ++++------ 9 files changed, 26 insertions(+), 39 deletions(-) diff --git a/benchmarks/run_xbow_benchmark.sh b/benchmarks/run_xbow_benchmark.sh index 9efaf76..c77c940 100755 --- a/benchmarks/run_xbow_benchmark.sh +++ b/benchmarks/run_xbow_benchmark.sh @@ -248,7 +248,7 @@ echo "[+] Launching eval agent with uv run" echo "[+] Logging uv run output to: $LOG_FILE" ( cd "$REPO_ROOT/deadend_cli/src/deadend_cli" && \ - uv run main.py eval-agent --eval-metadata-file "$META_FILE" --llm-providers azure_ai + uv run main.py eval-agent --eval-metadata-file "$META_FILE" --provider azure_ai --model-name Kimi-K2.5 ) 2>&1 | tee "$LOG_FILE" echo "[+] Stopping benchmark services with make stop" diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py index 6398c07..1887ca9 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py @@ -498,7 +498,7 @@ async def run( subtasks, website_info, exploit_info = await self.planner.expand( root, - context=context, + context=context if context else "", usage=RunUsage(), usage_limits=UsageLimits(request_limit=None), ) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/planner.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/planner.py index 446d94c..a4be820 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/planner.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/planner.py @@ -61,7 +61,7 @@ async def expand( context: str, usage: RunUsage, usage_limits: UsageLimits, - ) -> tuple[list[TaskNode], GeneralInfoOutput]: + ) -> tuple[list[TaskNode], GeneralInfoOutput, ExploitInfo]: """Expand a parent task into subtasks. Args: @@ -146,7 +146,7 @@ async def update_plan( context: str, usage: RunUsage, usage_limits: UsageLimits, - ) -> tuple[list[TaskNode], GeneralInfoOutput | ExploitInfo]: + ) -> tuple[list[TaskNode], GeneralInfoOutput | ExploitInfo, GeneralInfoOutput | ExploitInfo]: """Update and refine the plan for all tasks that share the same parent. This function takes a task node, finds all tasks that share the same parent diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py index c47d463..9a487e7 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py @@ -308,7 +308,7 @@ async def check( if not isinstance(judge_output, _JudgeOutput): logger.debug("JudgeAgentStrategy: unexpected output type %s", type(judge_output)) return ValidationVerdict(stop=False, confidence=0.0) - print(judge_output) + # print(judge_output) return ValidationVerdict( stop=judge_output.valid, confidence=judge_output.confidence_score, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py index 0d3aa3d..c65eeb5 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/factory.py @@ -373,7 +373,7 @@ def _extract_model_info(self, model: ModelSpec) -> tuple[str, str | None, str | return fq_model_name, api_key, api_base - def _create_fallback_output(self, error_msg: str, error_type: str) -> BaseModel: + def _create_fallback_output(self, error_msg: str, error_type: str) -> AgentOutput: """Create a fallback output with the correct schema type. Args: diff --git a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py index b76e172..f978358 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py @@ -444,7 +444,7 @@ async def threat_model(self, task: str): message_history="" ) - return task_node, context, validation_token + return task_node, threat_model_data, validation_token async def threat_model_stream(self, task: str): """Execute the threat modeling and orchestration workflow. diff --git a/deadend_cli/deadend_eval/src/deadend_eval/eval.py b/deadend_cli/deadend_eval/src/deadend_eval/eval.py index 1aafec9..a69723b 100644 --- a/deadend_cli/deadend_eval/src/deadend_eval/eval.py +++ b/deadend_cli/deadend_eval/src/deadend_eval/eval.py @@ -8,6 +8,7 @@ security research capabilities, and workflow effectiveness using various evaluation metrics and testing scenarios. """ +from deadend_agent.agents import AgentOutput from datetime import datetime from pathlib import Path @@ -87,17 +88,9 @@ async def eval_deadend_agent( eval_metadata: EvalMetadata, hard_prompt: bool, # choosing between hard and soft prompt - guided: bool, - # If guided enabled, the evaluation runs also on the subtasks - human_intervention: bool, - # whether or not ask user to specify information. - with_context_engine: bool, - # With context engineering enabled with_code_indexing: bool, # With code indexing enabled, code RAG specific to the application with_knowledge_base: bool, - # Knowledge base represents the database RAG added for notes or technical documents. - output_report: str ): """Evaluate an AI agent's performance on a security challenge. @@ -223,16 +216,11 @@ async def eval_deadend_agent( # Safely extract detailed_summary from threat_model_data # When supervisor fails, detailed_summary may not be set - detailed_summary = threat_model_data.get('detailed_summary', '') - if not detailed_summary: - # Fallback to last_output if available - last_output = threat_model_data.get('last_output', '') - if isinstance(last_output, dict): - detailed_summary = last_output.get('detailed_summary', str(last_output)) - else: - detailed_summary = str(last_output) if last_output else "No threat model data available" - - plan, validation_token = await deadend_agent.run_exploitation(threat_model=detailed_summary, task=prompt) + if isinstance(threat_model_data.output, AgentOutput): + detailed_summary = threat_model_data.output.get('detailed_summary', '') + else: + detailed_summary = threat_model_data.output + task_node, plan, validation_token = await deadend_agent.run_exploitation(threat_model=detailed_summary, task=prompt) if validation_token and len(validation_token) > 1 and eval_metadata.solution: solved = (validation_token == eval_metadata.solution) print(f"Validation check: expected={eval_metadata.solution}, got={validation_token}, solved={solved}") diff --git a/deadend_cli/src/deadend_cli/cli.py b/deadend_cli/src/deadend_cli/cli.py index 01a261e..c2a8f06 100644 --- a/deadend_cli/src/deadend_cli/cli.py +++ b/deadend_cli/src/deadend_cli/cli.py @@ -111,8 +111,11 @@ def eval_agent( None, help="Dataset file containing all the information about the challenges to run", ), - llm_providers: List[str] = typer.Option( - ["openai"], help="Specify the eval providers" + provider: str = typer.Option( + default="azure_ai", help="Provider name" + ), + model_name: str = typer.Option( + default="Kimi-K2.5", help="Model name" ), guided: bool = typer.Option( False, help="Run subtasks instead of one general task." @@ -150,8 +153,8 @@ def eval_agent( eval_interface( config=config, eval_metadata_file=eval_metadata_file, - providers=llm_providers, - guided=guided, + provider=provider, + model_name=model_name, ) ) finally: @@ -161,8 +164,8 @@ def eval_agent( @app.command() def init(): - """Initialize CLI config by prompting for env vars and saving to cache TOML. + """Initialize CLI config by prompting for env vars and saving to cache JSON. - Writes to ~/.cache/deadend/config.toml + Writes to ~/.cache/deadend/config.json """ init_cli_config() diff --git a/deadend_cli/src/deadend_cli/eval.py b/deadend_cli/src/deadend_cli/eval.py index 4bbcc30..cfc82e9 100644 --- a/deadend_cli/src/deadend_cli/eval.py +++ b/deadend_cli/src/deadend_cli/eval.py @@ -17,8 +17,8 @@ async def eval_interface( config: Config, eval_metadata_file: str, - providers: list[str], - guided: bool, + provider: str, + model_name: str, ): """Run evaluation interface for testing AI agent performance. @@ -66,7 +66,7 @@ async def eval_interface( model_registry = ModelRegistry(config=config) if not model_registry.has_any_model(): raise RuntimeError(f"No LM model configured. You can run `deadend init` to \ - initialize the required Model configuration for {providers[0]}") + initialize the required Model configuration for {provider}:{model_name}") # Initialize SQLite-based RAG rag_manager = init_rag_session_manager(storage_root=config.agents_storage_root) @@ -93,17 +93,13 @@ async def eval_interface( sandbox = sandbox_manager.get_sandbox(sandbox_id=sandbox_id) embedder_client = model_registry.get_embedder_model() await eval_deadend_agent( - model=model_registry.get_model(provider=providers[0]), + model=model_registry.get_model(provider=provider, model_name=model_name), embedder_client=embedder_client, code_indexer_db=rag_db, sandbox=sandbox, eval_metadata=eval_metadata, - guided=guided, - human_intervention=False, - with_context_engine=True, with_code_indexing=True, with_knowledge_base=True, - output_report="./", hard_prompt=False ) From eb79b6fe94b7c325884c488a9e1a885c8ba174d3 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Mon, 30 Mar 2026 16:21:18 +0200 Subject: [PATCH 09/12] The validation strategy rely now on just after the executor start. making it more suitable for usage. --- .../src/deadend_agent/agents/architecture.py | 109 ++++---------- .../agents/components/executor.py | 81 ++++++++++- .../components/validation_strategies.py | 41 ++++-- .../src/deadend_agent/deadend_agent.py | 134 ++++++++++++++---- .../deadend_eval/src/deadend_eval/eval.py | 23 +-- 5 files changed, 256 insertions(+), 132 deletions(-) diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py index 1887ca9..982297e 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/architecture.py @@ -1,14 +1,16 @@ from __future__ import annotations from typing import Any, Literal, AsyncGenerator -from uuid import UUID from pydantic_ai.usage import RunUsage, UsageLimits -from deadend_agent.agents.components.executor import AgentExecutor, ResultEvent, LogEvent +from deadend_agent.agents.components.executor import ( + AgentExecutor, + LogEvent, + ResultEvent, + ValidationStopEvent, +) from deadend_agent.agents.components.planner import Planner, TaskNode -from deadend_agent.agents.components.validation_strategies import ValidationGate, ValidationVerdict -from deadend_agent.agents.reporter import ReporterAgent from deadend_agent.context import ContextEngine from deadend_agent.logging import logger from deadend_agent.utils.structures import TaskPlanner @@ -42,9 +44,9 @@ class ADaPTAgent: This implementation follows the algorithm presented in the paper, by retaking the most relevant information: - - Usage of 3 components Executors, Planner and Validator. + - Usage of execution and planning components, with root-goal validation + handled centrally by the shared executor. """ - session_id: UUID task_node: TaskNode max_depth: int context: ContextEngine @@ -59,35 +61,22 @@ class ADaPTAgent: def __init__( self, - session_id: UUID, - agent_id: UUID, context: ContextEngine, executor: AgentExecutor, planner: Planner, - validation_gate: ValidationGate, - reporter: ReporterAgent, max_depth: int = 3, ): """Initialize the ADaPT agent. Args: - session_id: Unique identifier for this ADaPT session. - agent_id: Local agent ID used for AVFS workspace mounts. context: Shared context engine across all agents. executor: AgentExecutor instance for executing tasks. planner: Planner instance for decomposing tasks. - validation_gate: Composable gate that checks whether the - root goal is satisfied after every supervisor return. - reporter: ReporterAgent that writes an MD report on stop. max_depth: Maximum depth for task decomposition (default: 3). """ - self.session = session_id - self.agent_id = agent_id self.max_depth = max_depth self.executor = executor self.planner = planner - self.validation_gate = validation_gate - self.reporter = reporter self.context = context # Track attempted tasks to prevent redundant retries. @@ -102,10 +91,11 @@ async def _solve( ) -> AsyncGenerator[str | dict[str, Any], None]: """Recursively solve a task node using the ADaPT algorithm. - After every supervisor return the validation gate is consulted. - If the gate says stop (root goal achieved), write a report and - propagate exit_loop. Otherwise fall through to the ADaPT policy - (fail / expand / refine). + Each supervisor run is executed through the shared executor. The + executor owns root-goal validation and can emit a stop event when the + objective has been solved. ADaPT only reacts to that event and then + falls through to the normal policy (fail / expand / refine) when the + root goal is still unresolved. Args: node: The TaskNode to solve. @@ -210,6 +200,17 @@ def emit(message: str) -> str: ) break + elif isinstance(event, ValidationStopEvent): + node.status = "completed" + node.confidence_score = event.confidence_score + self.context.mark_task_completed(node.task, event.confidence_score) + yield emit( + f"[VALIDATION] Root goal achieved (confidence={event.confidence_score:.2f})" + ) + yield event + yield {"exit_loop": True} + return + elif isinstance(event, LogEvent): self.context.structured.append_to_log(event.message) yield emit(event.message) @@ -243,51 +244,12 @@ def emit(message: str) -> str: if detailed_summary: yield emit(f"[RESULT] {detailed_summary[:200]}") - # ----- 2. Validation gate ----- - # Only run the full gate (which may include expensive LLM judge) - # when there is reason to believe the root goal might be done: - # either the supervisor says the subtask is achieved, or - # confidence is high enough to warrant checking. - supervisor_output = { - "task_achieved": task_achieved, - "detailed_summary": detailed_summary, - "proofs": proofs, - "confidence_score": confidence_score, - } - - should_validate = ( - task_achieved - or confidence_score >= self.VALIDATE_THRESHOLD - ) - - if should_validate: - validation_context = self.context.get_unified_context(max_tokens=8000) - verdict = await self.validation_gate.check( - output=supervisor_output, - root_goal=self.context.final_goal, - context=validation_context, - ) - - if verdict.stop: - node.status = "completed" - node.confidence_score = verdict.confidence - self.context.mark_task_completed(node.task, verdict.confidence) - yield emit(f"[VALIDATION] Root goal achieved (confidence={verdict.confidence:.2f})") - - # Write report via the reporter agent. - await self._write_report(verdict) - - if verdict.token: - yield {"validation_token": verdict.token} - yield {"exit_loop": True} - return - - # ----- 3. If subtask done but root goal not yet, continue ----- + # ----- 2. If subtask done but root goal not yet, continue ----- if task_achieved: yield emit(f"[SUPERVISOR] Subtask completed: {node.task[:50]}...") return - # ----- 4. ADaPT policy (expand / refine / fail) ----- + # ----- 3. ADaPT policy (expand / refine / fail) ----- decision = self._policy(confidence_score) logger.debug( "task: %s decision: %s confidence: %.2f", @@ -448,24 +410,6 @@ def _policy(self, confidence_score: float) -> Literal["fail", "expand", "refine" return "validate" return "refine" - async def _write_report(self, verdict: ValidationVerdict) -> None: - """Delegate report generation and persistence to the ReporterAgent. - - The reporter agent has the write_workspace_file tool and writes the report - itself during execution — no programmatic write after the call. - Uses agent_id (not session_id) to match the mounted AVFS workspace. - """ - try: - await self.reporter.summarize_and_write( - root_goal=self.context.final_goal, - verdict=verdict, - context=self.context.get_unified_context(max_tokens=100_000), - session_id=str(self.agent_id), - ) - except Exception as exc: - # Report writing must never crash the agent loop. - logger.warning("ReporterAgent failed to write report: %s", exc) - async def run( self, task: str, @@ -494,7 +438,6 @@ async def run( parent=None, children=[], ) - self.context.set_root_task(root.task) subtasks, website_info, exploit_info = await self.planner.expand( root, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py index 4d2d6b2..89cee7f 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/executor.py @@ -12,6 +12,8 @@ MemoryAgent, ) from deadend_agent.agents.components.planner import TaskNode +from deadend_agent.agents.components.validation_strategies import ValidationGate, ValidationInput +from deadend_agent.agents.reporter import ReporterAgent from deadend_agent.context import ContextEngine from deadend_agent.config.settings import ModelSpec from deadend_agent.tools.avfs.write import write_text @@ -30,8 +32,18 @@ class ResultEvent(BaseModel): confidence_score: float context: dict[str, Any] + +class ValidationStopEvent(BaseModel): + """Event emitted when validation confirms the root objective is solved.""" + + type: Literal["validation_stop"] = "validation_stop" + validation_token: str = "" + confidence_score: float + critique: str = "" + reporter_output: str = "" + # Union type for all possible executor events -ExecutorEvent = LogEvent | ResultEvent +ExecutorEvent = LogEvent | ResultEvent | ValidationStopEvent def _memory_prompt_prefix(memory_context: str) -> str: @@ -92,7 +104,9 @@ def __init__( available_agents: dict[str, str] | None = None, agent_factory: Any | None = None, requires_approval: bool = False, - session_id: str | None = None + session_id: str | None = None, + validation_gate: ValidationGate | None = None, + reporter: ReporterAgent | None = None, ) -> None: """Initialize the AgentExecutor. @@ -109,6 +123,8 @@ def __init__( self.requires_approval = requires_approval self.context = context self.session_id = session_id + self.validation_gate = validation_gate + self.reporter = reporter self.memory_context = "" self.auth_session_key = "" @@ -145,6 +161,55 @@ def set_auth_session_key(self, auth_session_key: str) -> None: """Register the auth storage session key used by the python interpreter agent.""" self.auth_session_key = auth_session_key + def _build_validation_input( + self, + confidence_score: float, + result_context: dict[str, Any], + ) -> ValidationInput: + """Convert the latest supervisor result into structured validation input.""" + return ValidationInput( + task_achieved=result_context.get("task_achieved", False), + detailed_summary=result_context.get("detailed_summary", ""), + proofs=result_context.get("proofs", ""), + confidence_score=confidence_score, + ) + + async def _run_validation_and_report( + self, + validation_input: ValidationInput, + ) -> ValidationStopEvent | None: + """Validate the root goal and write the success report when solved. + + Validation is executed at the executor boundary so every supervisor + result is checked consistently, regardless of whether the caller is the + ADaPT planner or a direct top-level workflow entrypoint. + """ + if self.validation_gate is None or self.reporter is None: + return None + if not self.context.final_goal: + return None + + verdict = await self.validation_gate.check( + output=validation_input, + root_goal=self.context.final_goal, + context=self.context.get_unified_context(max_tokens=8000), + ) + if not verdict.stop: + return None + + reporter_output = await self.reporter.summarize_and_write( + root_goal=self.context.final_goal, + verdict=verdict, + context=self.context.get_unified_context(max_tokens=100_000), + session_id=str(self.session_id), + ) + return ValidationStopEvent( + validation_token=verdict.token, + confidence_score=verdict.confidence, + critique=verdict.critique, + reporter_output=reporter_output, + ) + async def _refresh_memory_context_for_task(self, task_query: str) -> str: """Retrieve task-specific memory immediately before supervisor execution.""" memory_workspace_root = ( @@ -219,7 +284,8 @@ async def execute_supervisor( Yields: LogEvent instances for streaming updates. - The final event is a ResultEvent instance. + The final event is either a ResultEvent for continued execution or + a ValidationStopEvent when the root task has been solved. """ context: dict[str, Any] = {"log": ""} @@ -516,6 +582,15 @@ async def call_memory_agent(ctx: RunContext[SupervisorDeps], prompt: str) -> str context["task_achieved"] = False context["proofs"] = "" + validation_input = self._build_validation_input( + confidence_score=confidence_score, + result_context=context, + ) + validation_event = await self._run_validation_and_report(validation_input) + if validation_event is not None: + yield validation_event + return + yield ResultEvent( confidence_score=confidence_score, context=context, diff --git a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py index 9a487e7..62505f8 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/agents/components/validation_strategies.py @@ -61,6 +61,20 @@ class ValidationVerdict(BaseModel): report: str = "" +class ValidationInput(BaseModel): + """Structured supervisor output used by root-goal validation strategies. + + Keeping validation input in a Pydantic model makes the contract between the + executor and the validation layer explicit and avoids passing loosely-shaped + dictionaries through the workflow. + """ + + task_achieved: bool = False + detailed_summary: str = "" + proofs: str = "" + confidence_score: float = 0.0 + + # --------------------------------------------------------------------------- # YAML-driven configuration # --------------------------------------------------------------------------- @@ -150,15 +164,14 @@ class ValidationStrategy(Protocol): async def check( self, - output: dict, + output: ValidationInput, root_goal: str, context: str, ) -> ValidationVerdict: """Evaluate whether *root_goal* is satisfied. Args: - output: The supervisor output dict (keys: task_achieved, - detailed_summary, proofs, confidence_score). + output: Structured supervisor output used for validation. root_goal: The top-level goal of the entire assessment. context: Accumulated execution context (unified context string). @@ -184,14 +197,14 @@ def __init__(self, pattern: str = r"FLAG\{[^}]+\}"): async def check( self, - output: dict, + output: ValidationInput, root_goal: str, context: str, ) -> ValidationVerdict: # Search in order: proofs first (most likely), then summary, then full context. searchable_fields = [ - output.get("proofs", ""), - output.get("detailed_summary", ""), + output.proofs, + output.detailed_summary, context, ] for field in searchable_fields: @@ -210,15 +223,15 @@ async def check( return ValidationVerdict(stop=False, confidence=0.0) @staticmethod - def _build_report(output: dict, token: str, root_goal: str) -> str: + def _build_report(output: ValidationInput, token: str, root_goal: str) -> str: return ( "# Validation Report — Flag Captured\n\n" f"**Goal:** {root_goal}\n\n" f"**Token:** `{token}`\n\n" "## Evidence\n\n" - f"```\n{output.get('proofs', 'N/A')}\n```\n\n" + f"```\n{output.proofs or 'N/A'}\n```\n\n" "## Summary\n\n" - f"{output.get('detailed_summary', 'N/A')}\n" + f"{output.detailed_summary or 'N/A'}\n" ) @@ -265,7 +278,7 @@ def __init__( async def check( self, - output: dict, + output: ValidationInput, root_goal: str, context: str, ) -> ValidationVerdict: @@ -282,9 +295,9 @@ async def check( "satisfied based on the execution trace.\n\n" f"## Goal\n{root_goal}\n\n" f"## Latest Supervisor Output\n" - f"Summary: {output.get('detailed_summary', '')}\n" - f"Proofs: {output.get('proofs', '')}\n" - f"Confidence: {output.get('confidence_score', 0.0)}\n\n" + f"Summary: {output.detailed_summary}\n" + f"Proofs: {output.proofs}\n" + f"Confidence: {output.confidence_score}\n\n" "# Execution Trace\n" f"{context}\n\n" "# Instructions\n" @@ -361,7 +374,7 @@ def __init__(self, strategies: list[ValidationStrategy]): async def check( self, - output: dict, + output: ValidationInput, root_goal: str, context: str, ) -> ValidationVerdict: diff --git a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py index f978358..a22c794 100644 --- a/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py +++ b/deadend_cli/deadend_agent/src/deadend_agent/deadend_agent.py @@ -5,6 +5,7 @@ from deadend_agent.logging import logger +from pydantic import BaseModel from pydantic_ai import RunUsage, UsageLimits from deadend_agent.config.settings import Config, ModelSpec from deadend_agent.models.registry import EmbedderClient @@ -15,10 +16,13 @@ from deadend_agent.agents.reporter import ReporterAgent, ReporterDeps from deadend_agent.agents.architecture import ADaPTAgent from deadend_agent.agents.generic_agents.memory_agent import MemoryAgent -from deadend_agent.agents.components.executor import AgentExecutor, ResultEvent +from deadend_agent.agents.components.executor import ( + AgentExecutor, + ResultEvent, + ValidationStopEvent, +) from deadend_agent.agents.components.planner import Planner, TaskNode from deadend_agent.agents.components.validation_strategies import ( - ValidationConfig, ValidationGate, build_validation_gate, load_validation_config, @@ -38,6 +42,21 @@ ApprovalCallback = Callable[..., Awaitable[str]] +class WorkflowStopResult(BaseModel): + """Stores the solved root-task result for the lifetime of a workflow run. + + Top-level workflow methods use this model to stop subsequent phases once a + validated solution has been found. The executor produces the validation + event, and DeadEndAgent persists the outcome so later phases can exit early. + """ + + solved: bool + validation_token: str = "" + confidence_score: float = 0.0 + critique: str = "" + reporter_output: str = "" + + class DeadEndAgent: """Main orchestrator for the DeadEnd security research framework.""" @@ -65,6 +84,7 @@ class DeadEndAgent: webapprecon_deps: WebappreconDeps | None = None challenge_name: str | None = None local_agent_id: UUID + stop_result: WorkflowStopResult | None = None def __init__( @@ -138,9 +158,35 @@ def reset_workflow_state(self) -> Generator[str, None, None]: """ self.goal_achieved = False self.interrupted = False + self.stop_result = None yield "[green]Workflow state reset for new execution[/green]" + def _create_completed_task_node(self, task: str) -> TaskNode: + """Create a completed task node for early-return success paths.""" + confidence_score = self.stop_result.confidence_score if self.stop_result else 1.0 + return TaskNode( + task=task, + depth=0, + confidence_score=confidence_score, + status="completed", + parent=None, + children=[], + ) + + def _record_validation_stop(self, event: ValidationStopEvent) -> WorkflowStopResult: + """Persist a validated root-task solution for the rest of the workflow run.""" + stop_result = WorkflowStopResult( + solved=True, + validation_token=event.validation_token, + confidence_score=event.confidence_score, + critique=event.critique, + reporter_output=event.reporter_output, + ) + self.goal_achieved = True + self.stop_result = stop_result + return stop_result + def set_approval_callback(self, callback): """Set a callback function for user approval input. @@ -332,7 +378,9 @@ def prepare_dependencies( model=self.model, context=self.context, available_agents=self.available_agents, - session_id=str(self.agent_id) + session_id=str(self.agent_id), + validation_gate=self.validation_gate, + reporter=self.reporter, ) self.executor.set_auth_session_key(self._target_session_key()) self.executor.set_memory_context(self.memory_context) @@ -355,11 +403,17 @@ async def threat_model(self, task: str): Returns: TaskNode with the execution plan and results """ + if self.goal_achieved and self.stop_result is not None: + return ( + self._create_completed_task_node(task), + self.stop_result.reporter_output, + self.stop_result.validation_token, + ) + # Simplified: directly use the executor's supervisor pattern # instead of going through the full ADaPT agent loop validation_token: str = "" - traces: list[str | dict[str, Any]] = [] prompt_task = f""" Prepare the necessary information (reconnaissance) to achieve the following task: {task} @@ -395,8 +449,6 @@ async def threat_model(self, task: str): # Set root task in context self.context.set_root_task(task) - # Get unified context for the executor - unified_context = self.context.get_unified_context(max_tokens=6000) target_context =f"Target : {self.context.target}" context = {} confidence_score = 0.0 @@ -407,7 +459,12 @@ async def threat_model(self, task: str): usage=RunUsage(), usage_limits=UsageLimits(request_limit=None, tool_calls_limit=None) ): - traces.append(event.model_dump()) + if isinstance(event, ValidationStopEvent): + stop_result = self._record_validation_stop(event) + task_node.confidence_score = event.confidence_score + task_node.status = "completed" + return task_node, stop_result.reporter_output, stop_result.validation_token + if isinstance(event, ResultEvent): confidence_score = event.confidence_score context = event.context @@ -456,6 +513,10 @@ async def threat_model_stream(self, task: str): Returns: TaskNode with the execution plan and results """ + if self.goal_achieved and self.stop_result is not None: + yield self.stop_result.reporter_output + return + prompt_task = f""" Prepare the necessary information (reconnaissance) to achieve the following task: {task} @@ -476,7 +537,6 @@ async def threat_model_stream(self, task: str): - Do NOT invent or guess endpoints - only use what is discovered - Return when you have gathered sufficient information to proceed with the task """ - validation_token: str = "" task_root = TaskNode( task=prompt_task, depth=0, @@ -486,7 +546,7 @@ async def threat_model_stream(self, task: str): children=[] ) - self.context.set_root_task(task_root.task) + self.context.set_root_task(task) target_context =f"Target : {self.context.target}" context = {} @@ -502,6 +562,12 @@ async def threat_model_stream(self, task: str): if self.interrupted: return # traces.append(event) + if isinstance(event, ValidationStopEvent): + stop_result = self._record_validation_stop(event) + task_root.confidence_score = event.confidence_score + task_root.status = "completed" + yield stop_result.reporter_output + return if isinstance(event, ResultEvent): confidence_score = event.confidence_score context = event.context @@ -543,6 +609,9 @@ async def threat_model_stream(self, task: str): async def run_exploitation(self, threat_model: str, task: str): """Runs the exploitation workflow""" + if self.goal_achieved and self.stop_result is not None: + return self._create_completed_task_node(task), self.stop_result.validation_token + # Create exploit agent as planner self.exploit_agent = PlannerExploitAgent( model=self.model, @@ -560,14 +629,11 @@ async def run_exploitation(self, threat_model: str, task: str): # Pass session_key as deps for the exploit agent self.planner = Planner(planner_agent=self.exploit_agent, deps=self._target_session_key()) + self.context.set_root_task(task) self.adapt_agent = ADaPTAgent( - session_id=self.session_id, - agent_id=self.agent_id, context=self.context, executor=self.executor, planner=self.planner, - validation_gate=self.validation_gate, - reporter=self.reporter, max_depth=self.max_depth, ) plan: TaskNode | None = None @@ -588,22 +654,28 @@ async def run_exploitation(self, threat_model: str, task: str): {threat_model} """ traces: list[str | dict[str, Any]] = [] - validation_token = "" async for event in self.adapt_agent.run(task=task_exploit): # interrupt signal if self.interrupted: return # Collect all events for trace saving - traces.append(event) + if hasattr(event, "model_dump"): + traces.append(event.model_dump()) + else: + traces.append(event) + + if isinstance(event, ValidationStopEvent): + stop_result = self._record_validation_stop(event) + if plan is None: + plan = self._create_completed_task_node(task) + return plan, stop_result.validation_token if isinstance(event, dict): if event.get("type") == "result": root_candidate = event.get("root") if isinstance(root_candidate, TaskNode): plan = root_candidate - elif event.get("validation_token"): - validation_token = event.get("validation_token") else: logger.debug("Event: %s", event.get("message", str(event))) else: @@ -612,10 +684,14 @@ async def run_exploitation(self, threat_model: str, task: str): if plan is None: raise RuntimeError("ADaPT agent did not produce a plan.") - return plan, validation_token + return plan, "" async def start_testing_stream(self, threat_model: str, task: str): """Runs the exploitation workflow""" + if self.goal_achieved and self.stop_result is not None: + yield self.stop_result.reporter_output + return + # Create exploit agent as planner self.exploit_agent = PlannerExploitAgent( model=self.model, @@ -624,17 +700,14 @@ async def start_testing_stream(self, threat_model: str, task: str): ) # We reset the context to make space and less confusion self.context.reset() + self.context.set_root_task(task) # Pass session_key as deps for the exploit agent self.planner = Planner(planner_agent=self.exploit_agent, deps=self._target_session_key()) self.adapt_agent = ADaPTAgent( - session_id=self.session_id, - agent_id=self.agent_id, context=self.context, executor=self.executor, planner=self.planner, - validation_gate=self.validation_gate, - reporter=self.reporter, max_depth=self.max_depth, ) plan: TaskNode | None = None @@ -648,6 +721,10 @@ async def start_testing_stream(self, threat_model: str, task: str): # interrupt signal if self.interrupted: return + if isinstance(event, ValidationStopEvent): + stop_result = self._record_validation_stop(event) + yield stop_result.reporter_output + return if isinstance(event, dict): if event.get("type") == "result": root_candidate = event.get("root") @@ -717,6 +794,10 @@ async def start_supervisor(self, task: str): Returns: TaskNode with the execution plan and results """ + if self.goal_achieved and self.stop_result is not None: + yield self.stop_result.reporter_output + return + prompt_task = f""" Your goal is to achieve the following task: {task} @@ -748,7 +829,6 @@ async def start_supervisor(self, task: str): ## The previous context if available is : {self.context.get_unified_context()} """ - validation_token: str = "" task_root = TaskNode( task=prompt_task, depth=0, @@ -758,7 +838,7 @@ async def start_supervisor(self, task: str): children=[] ) - self.context.set_root_task(task_root.task) + self.context.set_root_task(task) target_context =f"Target : {self.context.target}" context = {} @@ -770,6 +850,12 @@ async def start_supervisor(self, task: str): usage=RunUsage(), usage_limits=UsageLimits(request_limit=None, tool_calls_limit=None) ): + if isinstance(event, ValidationStopEvent): + stop_result = self._record_validation_stop(event) + task_root.confidence_score = event.confidence_score + task_root.status = "completed" + yield stop_result.reporter_output + return if isinstance(event, ResultEvent): confidence_score = event.confidence_score context = event.context diff --git a/deadend_cli/deadend_eval/src/deadend_eval/eval.py b/deadend_cli/deadend_eval/src/deadend_eval/eval.py index a69723b..4698d4c 100644 --- a/deadend_cli/deadend_eval/src/deadend_eval/eval.py +++ b/deadend_cli/deadend_eval/src/deadend_eval/eval.py @@ -202,6 +202,8 @@ async def eval_deadend_agent( if validation_token and len(validation_token) > 1 and eval_metadata.solution: solved = (validation_token == eval_metadata.solution) print(f"Validation check: expected={eval_metadata.solution}, got={validation_token}, solved={solved}") + elif deadend_agent.goal_achieved: + print(f"Validation stopped workflow in recon phase with token={validation_token}") else: print("Validation check: Continuing to the exploitation phase.") @@ -209,21 +211,27 @@ async def eval_deadend_agent( print(f"threat model is : {threat_model_data}") threat_model_computed = str(threat_model_data) - if not solved: + if not deadend_agent.goal_achieved: if len(validation_token) > 1: threat_model_computed += f"## Important NOTE\nThe flag found previously \ {str(validation_token)} is false and is not the right response. Find another way." - # Safely extract detailed_summary from threat_model_data - # When supervisor fails, detailed_summary may not be set - if isinstance(threat_model_data.output, AgentOutput): - detailed_summary = threat_model_data.output.get('detailed_summary', '') + # Pull the summary from the structured reporter output when available. + threat_model_output = getattr(threat_model_data, "output", threat_model_data) + if isinstance(threat_model_output, AgentOutput): + detailed_summary = threat_model_output.detailed_summary else: - detailed_summary = threat_model_data.output - task_node, plan, validation_token = await deadend_agent.run_exploitation(threat_model=detailed_summary, task=prompt) + detailed_summary = str(threat_model_output) + + task_node, validation_token = await deadend_agent.run_exploitation( + threat_model=detailed_summary, + task=prompt, + ) if validation_token and len(validation_token) > 1 and eval_metadata.solution: solved = (validation_token == eval_metadata.solution) print(f"Validation check: expected={eval_metadata.solution}, got={validation_token}, solved={solved}") + elif deadend_agent.goal_achieved: + print(f"Validation stopped workflow in exploitation phase with token={validation_token}") else: print("Validation check: FLAG NOT FOUND.") # Render and persist metrics for the end user. @@ -245,4 +253,3 @@ async def eval_deadend_agent( print(f"Deadend metrics summary written to {metrics_md_path}") print(f"Deadend metrics JSON written to {metrics_json_path}") print(metrics_md) - From c2267db11d2910d160adbddf3d9686b0956aa9f5 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Mon, 30 Mar 2026 21:49:07 +0200 Subject: [PATCH 10/12] reducing image size and rectifying the github workflow --- .github/workflows/docker-build.yml | 8 ++--- environments/images/kalilinux.Dockerfile | 37 ++++++++++++------------ 2 files changed, 22 insertions(+), 23 deletions(-) diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index f1d4b46..cedeb5c 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -4,12 +4,12 @@ on: push: branches: [ main, release_* ] paths: - - 'setup/images/kalilinux.Dockerfile' + - 'environments/images/kalilinux.Dockerfile' - '.github/workflows/docker-build.yml' pull_request: branches: [ main ] paths: - - 'setup/images/kalilinux.Dockerfile' + - 'environments/images/kalilinux.Dockerfile' - '.github/workflows/docker-build.yml' workflow_dispatch: inputs: @@ -20,7 +20,7 @@ on: env: REGISTRY: docker.io - IMAGE_NAME: bargacy/deadend-pentest + IMAGE_NAME: xoxruns/sandboxed_kali jobs: build-and-push: @@ -61,7 +61,7 @@ jobs: uses: docker/build-push-action@v5 with: context: . - file: ./setup/images/kalilinux.Dockerfile + file: ./environments/images/kalilinux.Dockerfile push: ${{ github.event_name != 'pull_request' }} tags: ${{ steps.meta.outputs.tags }} labels: ${{ steps.meta.outputs.labels }} diff --git a/environments/images/kalilinux.Dockerfile b/environments/images/kalilinux.Dockerfile index f1f6f47..156e65a 100644 --- a/environments/images/kalilinux.Dockerfile +++ b/environments/images/kalilinux.Dockerfile @@ -1,13 +1,19 @@ +# Build Go CLIs on the official Go image so the final image does not ship ~400MB+ of Go SDK. +# Matches linux/amd64 and linux/arm64 from buildx (golang image is multi-arch). +FROM golang:tip-trixie AS go-tools +RUN go install github.com/ffuf/ffuf@latest && \ + go install github.com/OJ/gobuster/v3@latest + FROM kalilinux/kali-rolling -ENV DEBIAN_FRONTEND=noninteractive -ENV PATH=$PATH:/usr/local/go/bin:/root/go/bin:/root/.cargo/bin:/root/.local/bin -ENV GOPATH=/root/go -ENV CARGO_HOME=/root/.cargo -ENV RUSTUP_HOME=/root/.rustup +ENV DEBIAN_FRONTEND=noninteractive \ + PATH=$PATH:/root/.local/bin \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 + +COPY --from=go-tools /go/bin/ffuf /go/bin/gobuster /usr/local/bin/ RUN apt-get update --fix-missing && \ - apt-get upgrade -y && \ apt-get install -y --no-install-recommends \ # Basic system utilities vim nano htop tree wget unzip tar gzip \ @@ -19,26 +25,19 @@ RUN apt-get update --fix-missing && \ python3 python3-pip python3-setuptools python3-dev python3-venv pipx \ # Security tools nmap masscan nikto dirb sqlmap hydra john hashcat \ + # Seclists + seclists \ # Additional security tools amap apt-utils bsdmainutils cewl crackmapexec crunch \ dnsenum dnsrecon dnsutils dos2unix enum4linux ftp hping3 \ joomscan kpcli libffi-dev mimikatz nasm nbtscan onesixtyone \ oscanner passing-the-hash patator php powershell powersploit \ theharvester whois wpscan && \ - # Install Go - wget -O go.tar.gz https://go.dev/dl/go1.21.5.linux-amd64.tar.gz && \ - tar -C /usr/local -xzf go.tar.gz && \ - rm go.tar.gz && \ - # Install Rust - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \ - # Install Go tools - /usr/local/go/bin/go install github.com/ffuf/ffuf@latest && \ - /usr/local/go/bin/go install github.com/OJ/gobuster/v3@latest && \ - # install semgrep pipx install semgrep && \ apt-get autoremove -y && \ apt-get clean && \ - rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \ - /usr/local/go/bin/go clean -cache -modcache -testcache + rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \ + /usr/share/doc/* /usr/share/man/* /usr/share/info/* \ + /root/.cache/pip /root/.cache/pipx -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] From 4ea92bb0a9063fdea6a8e1c2126dca2ef8187f25 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Mon, 30 Mar 2026 21:57:49 +0200 Subject: [PATCH 11/12] problem with powershell, completely removed --- environments/images/kalilinux.Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/environments/images/kalilinux.Dockerfile b/environments/images/kalilinux.Dockerfile index 156e65a..d8ceadb 100644 --- a/environments/images/kalilinux.Dockerfile +++ b/environments/images/kalilinux.Dockerfile @@ -31,8 +31,8 @@ RUN apt-get update --fix-missing && \ amap apt-utils bsdmainutils cewl crackmapexec crunch \ dnsenum dnsrecon dnsutils dos2unix enum4linux ftp hping3 \ joomscan kpcli libffi-dev mimikatz nasm nbtscan onesixtyone \ - oscanner passing-the-hash patator php powershell powersploit \ - theharvester whois wpscan && \ + oscanner passing-the-hash patator php\ + theharvester wpscan && \ pipx install semgrep && \ apt-get autoremove -y && \ apt-get clean && \ From abd1ea398b6770fe97e6cc6a48233b2f6cb378a0 Mon Sep 17 00:00:00 2001 From: Yassine Bargach Date: Mon, 30 Mar 2026 23:02:23 +0200 Subject: [PATCH 12/12] cleaning up the legacy chat written in python. was ugly af anyway. --- .github/workflows/docker-build.yml | 20 + deadend_cli/README.md | 68 +- deadend_cli/src/deadend_cli/banner.py | 41 -- deadend_cli/src/deadend_cli/chat.py | 834 ------------------------- deadend_cli/src/deadend_cli/cli.py | 65 +- deadend_cli/src/deadend_cli/console.py | 14 - 6 files changed, 24 insertions(+), 1018 deletions(-) delete mode 100644 deadend_cli/src/deadend_cli/banner.py delete mode 100644 deadend_cli/src/deadend_cli/chat.py delete mode 100644 deadend_cli/src/deadend_cli/console.py diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index cedeb5c..ac39588 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -1,3 +1,6 @@ +# Pushes to Docker Hub only on: push to main/release_* (not on pull_request), or workflow_dispatch. +# PRs build the image for CI but intentionally do not push (no registry credentials / avoid polluting tags). + name: Build and Push Deadend Docker Image on: @@ -36,6 +39,22 @@ jobs: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 + - name: Explain no push on pull requests + if: github.event_name == 'pull_request' + run: | + echo "PR workflow: image is built only (push=false). After merge to main or release_*, a push event will build and push to Docker Hub." + + - name: Require Docker Hub secrets when pushing + if: github.event_name != 'pull_request' + env: + DOCKER_USER: ${{ secrets.DOCKER_USERNAME }} + DOCKER_PASS: ${{ secrets.DOCKER_TOKEN }} + run: | + if [ -z "$DOCKER_USER" ] || [ -z "$DOCKER_PASS" ]; then + echo "::error::Add repository Actions secrets DOCKER_USERNAME and DOCKER_TOKEN (Docker Hub access token). Without them, the job cannot log in and nothing is pushed to Docker Hub." + exit 1 + fi + - name: Log in to Docker Hub if: github.event_name != 'pull_request' uses: docker/login-action@v3 @@ -58,6 +77,7 @@ jobs: type=raw,value=${{ github.event.inputs.tag }},enable=${{ github.event.inputs.tag != '' }} - name: Build and push Docker image + id: build uses: docker/build-push-action@v5 with: context: . diff --git a/deadend_cli/README.md b/deadend_cli/README.md index d1dc470..a7202d0 100644 --- a/deadend_cli/README.md +++ b/deadend_cli/README.md @@ -75,44 +75,6 @@ uv sync && uv build ```bash # Initialize configuration deadend-cli init - -# Start testing -deadend-cli chat \ - --target "http://localhost:3000" \ - --prompt "find SQL injection vulnerabilities" -``` - ---- - -## Usage Examples - -### Basic Vulnerability Testing - -```bash -# Test OWASP Juice Shop -docker run -p 3000:3000 bkimminich/juice-shop - -deadend-cli chat \ - --target "http://localhost:3000" \ - --prompt "test the login endpoint for SQL injection" -``` - -### API Security Testing - -```bash -deadend-cli chat \ - --target "https://api.example.com" \ - --prompt "test authentication endpoints" -``` - -### Autonomous Mode - -```bash -# Run without approval prompts (CTFs/labs only) -deadend-cli chat \ - --target "http://ctf.example.com" \ - --mode yolo \ - --prompt "find and exploit all vulnerabilities" ``` --- @@ -123,21 +85,13 @@ deadend-cli chat \ Initialize configuration and set up pgvector database -### `deadend-cli chat` - -Start interactive security testing session - -- `--target`: Target URL -- `--prompt`: Initial testing prompt -- `--mode`: `hacker` (approval required) or `yolo` (autonomous) - ### `deadend-cli eval-agent` Run evaluation against challenge datasets - `--eval-metadata-file`: Challenge dataset file -- `--llm-providers`: AI model providers to test -- `--guided`: Run with subtask decomposition +- `--provider`: AI model provider to use +- `--model-name`: AI model to use ### `deadend-cli version` @@ -287,24 +241,6 @@ Evaluated on XBOW's 104-challenge validation suite (black-box mode, January 2026 Strong performance: XSS (91%), Business Logic (86%), SQL injection (83%), IDOR (80%) Perfect scores: GraphQL, SSRF, NoSQL injection, HTTP method tampering (100%) ---- - -## Operating Modes - -**Hacker Mode (default):** Requires approval for dangerous operations - -```bash -deadend-cli chat --target URL --mode hacker -``` - -**YOLO Mode:** Autonomous execution (CTFs/labs only) - -```bash -deadend-cli chat --target URL --mode yolo -``` - ---- - ## Technology Stack - **LiteLLM**: Multi-provider model abstraction (OpenAI, Anthropic, Ollama) diff --git a/deadend_cli/src/deadend_cli/banner.py b/deadend_cli/src/deadend_cli/banner.py deleted file mode 100644 index c637925..0000000 --- a/deadend_cli/src/deadend_cli/banner.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (C) 2025 Yassine Bargach -# Licensed under the GNU Affero General Public License v3 -# See LICENSE file for full license information. - -"""Application banner and branding display module. - -This module provides the ASCII art banner and branding display functionality -for the security research CLI application, including rich text formatting -and visual presentation elements. -""" - -from rich import print -from rich.panel import Panel -from rich.box import ROUNDED - -from deadend_agent import Config - -BANNER = """[bold grey] -██████╗ ███████╗ █████╗ ██████╗ ███████╗███╗ ██╗██████╗ -██╔══██╗██╔════╝██╔══██╗██╔══██╗██╔════╝████╗ ██║██╔══██╗ -██║ ██║█████╗ ███████║██║ ██║█████╗ ██╔██╗ ██║██║ ██║ -██║ ██║██╔══╝ ██╔══██║██║ ██║██╔══╝ ██║╚██╗██║██║ ██║ -██████╔╝███████╗██║ ██║██████╔╝███████╗██║ ╚████║██████╔╝ -╚═════╝ ╚══════╝╚═╝ ╚═╝╚═════╝ ╚══════╝╚═╝ ╚═══╝╚═════╝ -[/bold grey] - -[bold yellow] PENETRATION TESTING CLI [/bold yellow] -[dim] Find vulnerabilities. Test defenses. Secure systems.[/dim] -""" - -def print_banner(config: Config): - print(BANNER) - print(Panel( - f"[bold]OpenAI API Key:[/bold] {'***' if config.openai_api_key else '[red]Not set[/red]'}\n" - f"[bold]DB URL:[/bold] {config.db_url or '[red]Not set[/red]'}\n" - f"[bold]Embedding Model:[/bold] {config.embedding_model}\n" - f"[bold]Log Level:[/bold] {config.log_level}", - title="Configuration", - border_style="blue", - box=ROUNDED - )) \ No newline at end of file diff --git a/deadend_cli/src/deadend_cli/chat.py b/deadend_cli/src/deadend_cli/chat.py deleted file mode 100644 index ee9b8ee..0000000 --- a/deadend_cli/src/deadend_cli/chat.py +++ /dev/null @@ -1,834 +0,0 @@ -# Copyright (C) 2025 Yassine Bargach -# Licensed under the GNU Affero General Public License v3 -# See LICENSE file for full license information. - -"""Interactive chat interface for security research and web application testing. - -This module provides a rich terminal-based chat interface that allows users to -interact with AI agents for security assessments, view real-time results, -and manage workflow execution through an intuitive conversational interface. -""" - -import time -from uuid import uuid4 -import asyncio -import sys -from enum import Enum -from typing import Dict, List, Callable, Optional -from deadend_agent.agents.recon_threatmodel_agent import ThreatModelOutput -from rich.console import Console -from rich.layout import Layout -from rich.panel import Panel -from rich.box import ROUNDED -from rich.table import Table -from rich import box -from rich.style import Style as RichStyle -from prompt_toolkit.application import Application -from prompt_toolkit.widgets import TextArea, Frame, Label, RadioList -from prompt_toolkit.layout import Layout as PTKLayout -from prompt_toolkit.key_binding import KeyBindings -from prompt_toolkit.styles import Style -from prompt_toolkit.layout.containers import HSplit -from prompt_toolkit.layout import Dimension as D -from pydantic import BaseModel -from pydantic_ai import DeferredToolRequests -from deadend_agent.utils.structures import Task -from deadend_agent.agents import RequesterOutput -from deadend_agent.agents.judge import JudgeOutput -from deadend_agent.utils.network import check_target_alive, deterministic_session_id -from deadend_agent import Config, DeadEndAgent, init_rag_session_manager, sandbox_setup, ModelRegistry -from deadend_agent.rag.session_manager import RagSessionManager -from .console import console_printer - -# Defining Agent modes -class Modes(str, Enum): - """CLI modes""" - yolo = "yolo" - hacker = "hacker" - -def print_pydantic_model(obj: BaseModel, title: str = "Agent Output") -> None: - """Print a Pydantic BaseModel object in a structured format. - - Args: - obj: The Pydantic BaseModel object to print - title: Title for the panel - """ - # Create a table to display the model fields - table = Table(show_header=True, header_style="bold magenta", box=box.ROUNDED) - table.add_column("Field", style="cyan", no_wrap=True) - table.add_column("Value", style="white") - - # Add each field to the table - for field_name, field_value in obj.model_dump().items(): - # Display the full value without truncation - display_value = str(field_value) - table.add_row(field_name, display_value) - - # Create a panel with the table - panel = Panel( - table, - title=f"[bold green]{title}[/bold green]", - border_style="green", - box=box.ROUNDED - ) - - console_printer.print(panel) - -class ChatInterface: - """Console chat interface utilities. - - Provides convenience helpers to render outputs with Rich and to prompt - for user input inside a panel using Prompt Toolkit. - """ - def __init__(self, max_history: int = 50): - self.console = Console() - self.max_history = max_history - self.conversation: List[Dict] = [] - self.total_tokens = 0 - self.layout = Layout() - self.layout.split_column( - Layout(name="header", size=4), - Layout(name="main", ratio=1), - Layout(name="footer", size=3) - ) - - async def wait_response(self, func: Callable, status: str, *args, **kwargs): - """Run an async function while showing a live status spinner. - - Args: - func: Awaitable/coroutine function to execute. - status: Status text to display alongside the spinner. - *args: Positional arguments forwarded to the function. - **kwargs: Keyword arguments forwarded to the function. - - Returns: - Any: The result returned by the awaited function. - """ - response = "" - start_time = time.time() - with self.console.status( - f"{status}", spinner="dots2" - ) as status_resp: - async def update_status(): - while True: - elapsed = time.time() - start_time - status_resp.update(f"{status} ({elapsed:.1f}s)") - await asyncio.sleep(0.1) - update_task = asyncio.create_task(update_status()) - try: - response = await func(*args, **kwargs) - return response - finally: - update_task.cancel() - try: - await update_task - except asyncio.CancelledError: - pass - - def print_chat_response(self, message: str, agent_name: str): - """Print a simple agent message to the console. - - Args: - message: Text content produced by the agent. - agent_name: Display name of the agent. - """ - # Panel(message, title=f"{agent_name}", border_style="magenta", box=ROUNDED) - self.console.print(f"{agent_name}: {message}") - - def print_requester_response(self, output: List[RequesterOutput], title: str): - """Render Requester outputs inside a Rich panel. - - Args: - output: Sequence of requester messages/responses. - title: Panel title. - """ - message = "" - for msg in output: - if msg.reasoning is not None: - text = f""" -[bold green]analysis :[/bold green] {msg.reasoning} -[bold green]status :[/bold green] {msg.state} -[bold green]raw response:[/bold green] - {msg.raw_response} - """ - else: - text = str(msg) - message += text - self.console.print(Panel(message, title=title, border_style="green", box=ROUNDED)) - - def print_planner_response(self, output: List[Task], title: str): - """Render a list of planner tasks inside a Rich panel. - - Args: - output: Tasks produced by the planner. - title: Panel title. - """ - message = "" - - for i in enumerate(output): - task=output[i] - # goal = task.goal - resp = task.output - # status = task.status - formatted_task = f""" -[bold magenta]step {i+1} :[/bold magenta]{resp} -""" - message += formatted_task - self.console.print(Panel(message, title=title, border_style="red", box=ROUNDED)) - - def startup(self): - """Print startup banner and basic usage help.""" - self.console.print("[bold green] Starting Agent: Chat interface.[/bold green]") - - async def ask_for_approval_panel(self, title: str = "Tool Approval Required") -> Optional[str]: - """Prompt for tool execution approval using Prompt Toolkit with choices. - - Args: - title: Title text to display with the confirmation dialog. - - Returns: - 'yes' for approval, 'no' for rejection, or None if cancelled - """ - try: - style = Style.from_dict({ - "frame.border": "ansiyellow", - "radio-checked": "ansigreen", - "radio-unchecked": "ansired", - }) - - # Define approval choices - choices = [ - ("yes", "Approve tool execution"), - ("no", "Deny tool execution"), - ] - - # Create radio list for choices, default to yes - radio_list = RadioList(choices, default="yes") - # Set default selection to "yes" - radio_list.current_value = "yes" - - # Create confirmation prompt text - prompt_text = Label(text="Do you approve these tool executions?", style="ansiwhite") - - # Create footer with instructions - footer_text = "Commands: ↑/↓=Navigate | Space=Toggle | Enter=Submit | Ctrl+C=Cancel" - footer = Label(text=footer_text, style="ansiblack") - - root_container = HSplit([ - Label(text=title, style="bold ansiyellow"), - prompt_text, - Frame(body=radio_list), - footer, - ]) - - kb = KeyBindings() - - @kb.add("c-c") - def _(event): # type: ignore[no-redef] - event.app.exit(result=None) # Cancel without selection - - # Override RadioList's Enter behavior to exit with result - def custom_enter_handler(event): - radio_list._handle_enter() - selected_value = radio_list.current_value - event.app.exit(result=selected_value) - - # Replace RadioList's enter binding - radio_list.control.key_bindings.add("enter")(custom_enter_handler) - - app = Application( - layout=PTKLayout(container=root_container), - key_bindings=kb, - full_screen=False, - mouse_support=False, - style=style, - min_redraw_interval=0.01, - ) - app.layout.focus(radio_list) - - return await app.run_async() - except KeyboardInterrupt: - console_printer.print("\n[yellow]Approval cancelled.[/yellow]") - return None - - async def ask_with_ptk_panel( - self, - title: str = "Prompt", - placeholder: str = "Type here and press Enter", - interrupt_callback: Callable | None = None - ) -> str: - """Prompt for input inside a bordered frame. - - Typing occurs inside the frame; Enter accepts, Esc/Ctrl-C cancels. - - Args: - title: Title text rendered above the frame. - placeholder: Prompt text shown before the cursor inside the field. - - Returns: - The string entered by the user, or None if cancelled. - """ - try: - style = Style.from_dict({ - "frame.border": "#696969", - "text-area": "", - }) - - input_field = TextArea( - multiline=True, - wrap_lines=True, - prompt=placeholder, - scrollbar=True, - text="", - height=D(min=1, max=5), - ) - - # Create footer with available commands - footer_text = "Commands: Ctrl+C=Exit | Ctrl+I=Interrupt | Enter=Submit | /help=Help | /clear=Clear | /new-target=New Target" - footer = Label(text=footer_text, style="ansiblack") - - root_container = HSplit([ - Label(text=title, style="bold #696969"), - Frame(body=input_field), - footer, - ]) - - kb = KeyBindings() - - def _accept_handler(_buff): - app.exit(result=input_field.text) - input_field.accept_handler = _accept_handler - - @kb.add("c-c") - def _(event): - event.app.exit(result=None) - - @kb.add("c-i") - def _(event): - if interrupt_callback: - interrupt_callback() - event.app.exit(result="__INTERRUPT__") - - @kb.add("enter") - def _(event): - text = input_field.text - if text.startswith("/"): - # Handle commands - if text == "/clear": - event.app.exit(result="__CLEAR__") - elif text == "/new-target": - event.app.exit(result="__NEW_TARGET__") - elif text == "/help": - event.app.exit(result="__HELP__") - elif text == "/quit": - event.app.exit(result=None) - else: - # Unknown command, treat as regular text - event.app.exit(result=text) - else: - event.app.exit(result=text) - - app = Application( - layout=PTKLayout(container=root_container), - key_bindings=kb, - full_screen=False, - mouse_support=False, - style=style, - min_redraw_interval=0.01, - ) - app.layout.focus(input_field) - - return await app.run_async() - except KeyboardInterrupt: - return "" - -async def chat_interface( - config: Config, - # Config - prompt: str, - # CLI user prompt - mode: Modes, - # Mode setup - target: str, - # target - openapi_spec, - # OpenAPI spec if available - knowledge_base: str, - # Knowledge base path - workspace_root: str | None = None, - # Optional AVFS workspace root - llm_provider: str = "openai" - # LLM provider - ): - """Chat Interface for the CLI""" - model_registry = ModelRegistry(config=config) - if not model_registry.has_any_model(): - raise RuntimeError(f"No LM model configured. You can run `deadend init` to \ - initialize the required Model configuration for {llm_provider}") - - model = model_registry.get_model(provider=llm_provider) - embedder_client = model_registry.get_embedder_model() - - # Initialize SQLite-based RAG session manager - rag_manager = init_rag_session_manager(storage_root=config.agents_storage_root) - local_agent_id = config.get_local_agent_id() - # Settings up sandbox - try: - sandbox_manager = sandbox_setup() - sandbox_id = sandbox_manager.create_sandbox("xoxruns/sandboxed_kali", network_name="host") - sandbox = sandbox_manager.get_sandbox(sandbox_id=sandbox_id) - except Exception as e: - console_printer.print(f"[yellow]Sandbox manager could not be started : {e}. Continuing without sandbox.[/yellow]") - - chat_interface = ChatInterface() - chat_interface.startup() - chat_interface.console.print(f"Model currently used : {model.model_name}") - user_prompt = prompt - - # Setup available agents - available_agents = { - 'requester': "Agent specialized in fine-grained testing and sending raw request data. Capable of handling authentication (session and token). Uses pupeteer in the background. Capable of exploring APIs and websites. Best for gathering auth tokens, testing individual endpoints, and precise request manipulation. Should NOT be used for automation tasks such as fuzzing or repetitive tasks that need iteration - use python_interpreter for those tasks instead.", - 'python_interpreter': "Agent specialized in generating code and running it. Each code generated is ran safely in a sandboxed webassembly. Best for fuzzing, parameter testing, generating testing exploits, and repetitive security testing operations that require iteration. Use this agent for tasks that need automation, loops, or multiple iterations.", - 'shell': "Agent that gives access to a terminal bash shell. Run linux commands here.", - 'memory': "Agent specialized in reading and writing the persistent memory workspace under the agent cache.", - 'router_agent': 'Router agent, expert that routes to the specific agent needed to achieve the next step of the plan.' - } - - - # Set up approval callback to use Prompt Toolkit - async def approval_callback(): - return await chat_interface.ask_for_approval_panel("Tool Execution Approval Required") - - # { - # 'webapp_recon': "Expert cybersecurity agent that enumerates a web \ - # target to understand the architecture and understand the endpoints\ - # and where an attack vector could be tested.", - # 'recon_shell': "Expert system reconnaissance agent that performs \ - # infrastructure-level security assessment using specialized \ - # command-line tools. Focuses on network scanning, system enumeration,\ - # and infrastructure analysis. Should NOT be used for simple request \ - # injections or basic web testing - use specialized web tools instead.", - # # 'planner_agent': 'Expert cybersecurity agent that plans what is the next step to do', - # 'router_agent': 'Router agent, expert that routes to the specific \ - # agent needed to achieve the next step of the plan.', - # 'python_interpreter_agent': 'Generate code and runs it in a python \ - # interpreter depending on the goal given.' - # } - # workflow_agent.register_agents(available_agents) - # Check if the provided target is reachable before proceeding - alive = False - while not alive: - if not target: - # Prompt user for target if none provided - console_printer.print("[yellow]No target specified. \ -Please provide a target URL.[/yellow]") - target = await chat_interface.ask_with_ptk_panel( - title="Target URL", - placeholder="Enter the target URL (e.g., https://example.com) > " - ) - - if not target: - console_printer.print("[red]No target provided. Exiting application...[/red]") - sys.exit(1) - - # Check target reachability - alive, status_code, err = await check_target_alive(target) - if alive: - console_printer.print(f"[green]Target reachable[/green] (status={status_code})") - else: - console_printer.print( - f"[red]Target not reachable[/red] (status={status_code}, error={err})") - console_printer.print("[yellow]Please provide a different target URL.[/yellow]") - target = await chat_interface.ask_with_ptk_panel( - title="Target URL", - placeholder="Enter a valid target URL (e.g., https://example.com) > " - ) - - if not target: - console_printer.print("[red]No target provided. Exiting application...[/red]") - sys.exit(1) - - # Create agent with a runtime session plus deterministic target embedding session - runtime_session_id = uuid4() - embedding_session_id = deterministic_session_id(target) - deadend_agent = DeadEndAgent( - session_id=runtime_session_id, - embedding_session_id=embedding_session_id, - model=model, - available_agents=available_agents, - max_depth=3, - workspace_root=workspace_root, - agents_storage_root=config.agents_storage_root, - local_agent_id=local_agent_id, - ) - deadend_agent.set_approval_callback(approval_callback) - - # Indexing webtarget - deadend_agent.init_webtarget_indexer(target=target) - web_ressource_crawl = await chat_interface.wait_response( - func=deadend_agent.crawl_target, - status="Gathering webpage resources..." - ) - code_chunks, embed_diff = await chat_interface.wait_response( - func=deadend_agent.embed_target, - status="Indexing the different webpage resources...", - embedder_client=embedder_client - - ) - if embed_diff: - changed = len(embed_diff.get("changed_files", [])) - removed = len(embed_diff.get("removed_files", [])) - console_printer.print( - f"[blue]Embedding diff[/blue] changed={changed} removed={removed} " - f"(session={deadend_agent.embedding_session_id})" - ) - - # Get per-session SQLite RAG connector - rag_db = await rag_manager.get_connector( - agent_id=local_agent_id, - embedding_session_id=embedding_session_id, - target=target, - ) - - # Sync embeddings to SQLite - if embed_diff: - delete_files = embed_diff.get("changed_files", []) + embed_diff.get("removed_files", []) - if delete_files: - await rag_db.delete_code_chunks_for_files(files=delete_files) - if code_chunks: - await chat_interface.wait_response( - func=rag_db.batch_insert_code_chunks, - status="Syncing DB", - code_chunks_data=code_chunks - ) - - # Preparing deadend Dependencies - deadend_agent.prepare_dependencies( - embedder_client=embedder_client, - rag_connector=rag_db, - sandbox=sandbox, - target=target - ) - - # Setup the knowledge base in the database if necessary - ## TODO: adding the knowledge base handler -# if knowledge_base: -# if os.path.exists(knowledge_base) and os.path.isdir(knowledge_base): -# workflow_agent.knowledge_base_init(folder_path=knowledge_base) -# kb_chunks = await chat_interface.wait_response( -# func=workflow_agent.knowledge_base_index, -# status="Indexing the knowledge base...", -# ) -# # insert to db -# insert_kn = await chat_interface.wait_response( -# func=rag_db.batch_insert_kb_chunks, -# status="Syncing DB", -# knowledge_chunks_data=kb_chunks -# ) -# else: -# console_printer.print(f"[yellow]Warning: Knowledge base folder '{knowledge_base}' \ -# does not exist or is not a directory. Skipping knowledge base initialization.[/yellow]") - - - # Agent interruption flag - agent_interrupted = False - - def interrupt_agent(): - nonlocal agent_interrupted - agent_interrupted = True - deadend_agent.interrupt_workflow() - console_printer.print("\n[yellow]Agent interrupted by user (Ctrl+I)[/yellow]") - - try: - while True: - # Check if user prompt is provided, ask for it if not - while not user_prompt: - user_prompt = await chat_interface.ask_with_ptk_panel( - title="User Prompt", - placeholder=">>> ", - interrupt_callback=interrupt_agent - ) - - if not user_prompt: - console_printer.print("[red]No prompt provided. Exiting...[/red]") - break - elif user_prompt == "__CLEAR__": - console_printer.print("[green]Context cleared[/green]") - # Clear conversation history - chat_interface.conversation = [] - user_prompt = None - continue - elif user_prompt == "__NEW_TARGET__": - # Get new target - new_target = await chat_interface.ask_with_ptk_panel( - title="New Target URL", - placeholder="Enter the new target URL (e.g., https://example.com) > " - ) - if new_target and new_target != "__CLEAR__" and new_target != "__NEW_TARGET__": - target = new_target - console_printer.print(f"[green]Target changed to: {target}[/green]") - # Recreate agent with deterministic session for new target - runtime_session_id = uuid4() - embedding_session_id = deterministic_session_id(target) - deadend_agent = DeadEndAgent( - session_id=runtime_session_id, - embedding_session_id=embedding_session_id, - model=model, - available_agents=available_agents, - max_depth=3, - workspace_root=workspace_root, - agents_storage_root=config.agents_storage_root, - local_agent_id=local_agent_id, - ) - deadend_agent.set_approval_callback(approval_callback) - # Re-initialize with new target - deadend_agent.init_webtarget_indexer(target=target) - web_ressource_crawl = await chat_interface.wait_response( - func=deadend_agent.crawl_target, - status="Gathering webpage resources for new target..." - ) - code_chunks, embed_diff = await chat_interface.wait_response( - func=deadend_agent.embed_target, - status="Indexing the different webpage resources for new target...", - embedder_client=embedder_client - ) - if embed_diff: - changed = len(embed_diff.get("changed_files", [])) - removed = len(embed_diff.get("removed_files", [])) - console_printer.print( - f"[blue]Embedding diff[/blue] changed={changed} removed={removed} " - f"(session={deadend_agent.embedding_session_id})" - ) - # Get new per-session SQLite connector for new target - rag_db = await rag_manager.get_connector( - agent_id=local_agent_id, - embedding_session_id=embedding_session_id, - target=target, - ) - if embed_diff: - delete_files = embed_diff.get("changed_files", []) + embed_diff.get("removed_files", []) - if delete_files: - await rag_db.delete_code_chunks_for_files(files=delete_files) - if code_chunks: - insert = await chat_interface.wait_response( - func=rag_db.batch_insert_code_chunks, - status="Syncing DB with new target data", - code_chunks_data=code_chunks - ) - # Re-prepare deps with the new connector - deadend_agent.prepare_dependencies( - embedder_client=embedder_client, - rag_connector=rag_db, - sandbox=sandbox, - target=target, - ) - user_prompt = None - continue - elif user_prompt == "__HELP__": - console_printer.print(""" -[bold cyan]Available Commands:[/bold cyan] - [bold]/help[/bold] - Show this help message - [bold]/clear[/bold] - Clear conversation context - [bold]/new-target[/bold] - Change the target URL - [bold]/quit[/bold] - Exit the application - -[bold cyan]Keyboard Shortcuts:[/bold cyan] - [bold]Ctrl+C[/bold] - Exit the application - [bold]Ctrl+I[/bold] - Interrupt running agent - [bold]Enter[/bold] - Submit input - """) - user_prompt = None - continue - elif user_prompt == "__INTERRUPT__": - console_printer.print("[yellow]Interrupt command received[/yellow]") - user_prompt = None - continue - elif user_prompt == "": - console_printer.print("[red]No prompt provided. Please try again.[/red]") - user_prompt = None - continue - - if not user_prompt: - break - - # Reset interruption flag and workflow state for new execution - agent_interrupted = False - deadend_agent.reset_workflow_state() - - judge_output = None - - threat_model = "" - try: - async for item in deadend_agent.threat_model_stream( - task=user_prompt - ): - # Skip printing DeferredToolRequests objects - if hasattr(item, 'output') and isinstance(item.output, DeferredToolRequests): - continue - if isinstance(item, ThreatModelOutput): - console_printer.print(f"[bold yellow]Target extracted information:[/bold yellow] {item.website_general_information}") - console_printer.print(f"[bold yellow]The tech stack is :[/bold yellow] {item.technology_stack}") - console_printer.print(f"[bold yellow]Discovered endpoints:[/bold yellow] {item.endpoints}") - threat_model += str(item.model_dump()) - if isinstance(item, str) and len(item) > 50: - console_printer.print(f"The threat model analyzed is : \n{item}") - threat_model += item - # Special handling for RequesterOutput - print just the reasoning - if hasattr(item, 'output') and isinstance(item.output, RequesterOutput): - console_printer.print(f"[bold green]Requester Analysis:[/bold green] {item.output.reasoning}") - console_printer.print(f"[bold green]Raw response:[/bold green] {item.output.raw_response}") - continue - - # Check if this is the final result (JudgeOutput) - if isinstance(item, JudgeOutput): - judge_output = item - - # Check if this is a Pydantic BaseModel object - if isinstance(item, BaseModel): - threat_model += str(item.model_dump()) - # Special handling for RouterOutput - print as simple text - if type(item).__name__ == "RouterOutput": - console_printer.print(f"[cyan]Router:[/cyan] \ -{item.next_agent_name}") - console_printer.print(f"[cyan]Reasoning:[/cyan] {item.reasoning}") - else: - # Determine the type of model for better title - model_type = type(item).__name__ - print_pydantic_model(item, f"{model_type} Output") - - elif isinstance(item, list) and len(item) > 0 and hasattr(item[0], 'goal'): - tasks_text = "" - for i, task in enumerate(item, 1): - tasks_text += f"[cyan]Step {i}:[/cyan]" - tasks_text += f"[white]{task.goal}[/white]\n" - tasks_text += f"[grey39]{task.output}[/grey39]\n" - - style = RichStyle(italic=True) - - task_panel = Panel( - tasks_text.strip(), - style=style, - title="Planned tasks", - border_style="grey39", - box=box.ROUNDED - ) - console_printer.print(task_panel) - else: - # Print regular string messages - console_printer.print(f"[bold red]No output format : [/bold red]{item}") - threat_model += str(item) - - # Check for interruption - if agent_interrupted or deadend_agent.interrupted: - break - # Small delay to allow for interruption - await asyncio.sleep(0.1) - except InterruptedError as e: - console_printer.print(f"[yellow]Workflow interrupted: {e}[/yellow]") - judge_output = None - except Exception as e: - console_printer.print(f"[red]Workflow error: {e}[/red]") - judge_output = None - - try: - async for item in deadend_agent.start_testing_stream( - task=user_prompt, - threat_model=threat_model - ): - # Skip printing DeferredToolRequests objects - if hasattr(item, 'output') and isinstance(item.output, DeferredToolRequests): - continue - if isinstance(item, ThreatModelOutput): - console_printer.print(f"[bold yellow]Target extracted information:[/bold yellow] {item.website_general_information}") - console_printer.print(f"[bold yellow]The tech stack is :[/bold yellow] {item.technology_stack}") - console_printer.print(f"[bold yellow]Discovered endpoints:[/bold yellow] {item.endpoints}") - if isinstance(item, str) and len(item) > 50: - console_printer.print(f"The threat model analyzed is : \n{item}") - # Special handling for RequesterOutput - print just the reasoning - if hasattr(item, 'output') and isinstance(item.output, RequesterOutput): - console_printer.print(f"[bold green]Requester Analysis:[/bold green] {item.output.reasoning}") - console_printer.print(f"[bold green]Raw response:[/bold green] {item.output.raw_response}") - continue - - # Check if this is the final result (JudgeOutput) - if isinstance(item, JudgeOutput): - judge_output = item - - # Check if this is a Pydantic BaseModel object - if isinstance(item, BaseModel): - # Special handling for RouterOutput - print as simple text - if type(item).__name__ == "RouterOutput": - console_printer.print(f"[cyan]Router:[/cyan] \ -{item.next_agent_name}") - console_printer.print(f"[cyan]Reasoning:[/cyan] {item.reasoning}") - else: - # Determine the type of model for better title - model_type = type(item).__name__ - print_pydantic_model(item, f"{model_type} Output") - - elif isinstance(item, list) and len(item) > 0 and hasattr(item[0], 'goal'): - tasks_text = "" - for i, task in enumerate(item, 1): - tasks_text += f"[cyan]Step {i}:[/cyan]" - tasks_text += f"[white]{task.goal}[/white]\n" - tasks_text += f"[grey39]{task.output}[/grey39]\n" - - style = RichStyle(italic=True) - - task_panel = Panel( - tasks_text.strip(), - style=style, - title="Planned tasks", - border_style="grey39", - box=box.ROUNDED - ) - console_printer.print(task_panel) - else: - # Print regular string messages - console_printer.print(item) - - # Check for interruption - if agent_interrupted or deadend_agent.interrupted: - break - - # Small delay to allow for interruption - await asyncio.sleep(0.1) - except InterruptedError as e: - console_printer.print(f"[yellow]Workflow interrupted: {e}[/yellow]") - judge_output = None - except Exception as e: - console_printer.print(f"[red]Workflow error: {e}[/red]") - judge_output = None - - - # Check if agent was interrupted - if agent_interrupted or deadend_agent.interrupted: - console_printer.print("[yellow]Agent execution was interrupted[/yellow]") - # Reset the workflow interruption flag for next execution - deadend_agent.interrupted = False - user_prompt = None - continue - - # Print judge output in a nice format - console_printer.print("[bold blue]Agent task completed[/bold blue]") - - if hasattr(judge_output, 'output') and hasattr(judge_output.output, 'goal_achieved'): - if judge_output.output.goal_achieved: - console_printer.print("[bold green]✓ Goal achieved[/bold green]") - else: - console_printer.print("[bold red]✗ Goal not achieved[/bold red]") - - if hasattr(judge_output, 'output') and hasattr(judge_output.output, 'reasoning'): - console_printer.print("\n[bold yellow]Reasoning:[/bold yellow]") - console_printer.print(f"{judge_output.output.reasoning}") - - if hasattr(judge_output, 'output') and hasattr(judge_output.output, 'solution'): - console_printer.print("\n[bold cyan]Solution:[/bold cyan]") - console_printer.print(f"{judge_output.output.solution}") - - user_prompt = None - - except KeyboardInterrupt: - console_printer.print("\n[yellow]Received Ctrl+C. Exiting gracefully...[/yellow]") - console_printer.print("[green]Thank you for using Deadend CLI![/green]") - sys.exit(0) diff --git a/deadend_cli/src/deadend_cli/cli.py b/deadend_cli/src/deadend_cli/cli.py index c2a8f06..35a69fd 100644 --- a/deadend_cli/src/deadend_cli/cli.py +++ b/deadend_cli/src/deadend_cli/cli.py @@ -4,12 +4,11 @@ """Deadend CLI entrypoint using Typer. -Defines commands to run interactive chat and evaluation agents. +Defines maintenance and evaluation commands for the Python package. """ import asyncio import importlib.metadata import os -from typing import List import logging import docker import typer @@ -17,10 +16,8 @@ from rich.console import Console from deadend_agent import config_setup from deadend_agent.core import start_python_sandbox -from .banner import print_banner from .cli_logging import setup_logging from .init import init_cli_config, check_docker -from .chat import Modes, chat_interface from .eval import eval_interface # Fix Docker socket path if default doesn't exist @@ -32,7 +29,7 @@ console = Console() -app = typer.Typer(help="Deadend CLI - interact with the Deadend framework.") +app = typer.Typer(help="Deadend CLI - Python maintenance and evaluation commands.") @app.command() @@ -47,64 +44,6 @@ def version(): ) -@app.command() -def chat( - prompt: str = typer.Option(None, help="Send a prompt directly to chat mode."), - target: str = typer.Option(None, help="Target URL or identifier for chat."), - mode: Modes = typer.Option( - Modes.hacker, help="Two modes available, yolo and hacker." - ), - openapi_spec: str = typer.Option( - None, help="Path to the OpenAPI specification file." - ), - knowledge_base: str = typer.Option(None, help="Folder path to the knowledge base."), - workspace_root: str = typer.Option(None, help="Host workspace to mount into AVFS."), -): - """Run the interactive chat agent. - - Args: - prompt: Optional initial prompt to pre-fill the chat. - target: Target host or URL context for the agent. - openapi_spec: Path to an OpenAPI spec to load for context. - """ - # Check Docker availability first - docker_client = docker.from_env() - if not check_docker(docker_client): - console.print( - "\n[red]Docker is required for this application to function properly.[/red]" - ) - console.print("Please install Docker from: https://docs.docker.com/get-docker/") - console.print( - "Make sure Docker daemon is running, then run this command again." - ) - raise typer.Exit(1) - - # Init configuration - config = config_setup() - log_level_name = str(config.log_level or "INFO").upper() - log_level = getattr(logging, log_level_name, logging.INFO) - setup_logging(level=log_level) - print_banner(config=config) - - python_process = start_python_sandbox() - console.print(f"Python sandbox started: {python_process}") - try: - asyncio.run( - chat_interface( - config=config, - prompt=prompt, - mode=mode, - target=target, - openapi_spec=openapi_spec, - knowledge_base=knowledge_base, - workspace_root=workspace_root, - ) - ) - finally: - if python_process.poll() is None: - python_process.terminate() - - @app.command() def eval_agent( eval_metadata_file: str = typer.Option( diff --git a/deadend_cli/src/deadend_cli/console.py b/deadend_cli/src/deadend_cli/console.py deleted file mode 100644 index 63c2113..0000000 --- a/deadend_cli/src/deadend_cli/console.py +++ /dev/null @@ -1,14 +0,0 @@ -# Copyright (C) 2025 Yassine Bargach -# Licensed under the GNU Affero General Public License v3 -# See LICENSE file for full license information. - -"""Console output management for the security research framework. - -This module provides centralized console output management using Rich -framework, ensuring consistent formatting and styling across the application -for better user experience and readability. -""" - -from rich.console import Console - -console_printer = Console() \ No newline at end of file