diff --git a/docs/eiq-agent-architecture.md b/docs/eiq-agent-architecture.md new file mode 100644 index 000000000..8c352dc04 --- /dev/null +++ b/docs/eiq-agent-architecture.md @@ -0,0 +1,1020 @@ +# Exploit Intelligence Agent Architecture + +## Overview + +The ExploitIQ CVE agent is a multi-stage investigation system that assesses whether flagged vulnerabilities are exploitable in a containerized application. The pipeline follows a **fan-out / fan-in** pattern: for each CVE, a checklist of investigation questions is generated upstream, then the `cve_agent` orchestrator processes every question concurrently. Each question is classified by an LLM dispatcher and routed to a specialized sub-agent. The sub-agent runs an autonomous **thought → tool → observation** loop until it can answer the question or exhausts its iteration budget. Results are collected and normalized into `checklist_results` for downstream summarization. + +At a high level, the process has two levels: + +1. **Checklist loop** - the dispatcher analyzes each question, routes it to a specialized sub-agent, and repeats until all checklist items are answered. +2. **Sub-agent state machine** - inside each routed question, a five-node loop (PreProcess → Thought → Tool → Observation) autonomously investigates using code analysis tools and structured reasoning until the question is answered or the iteration limit is reached. + +--- + +## Process Steps + +### 1. Pre-Process (Orchestration) + +**Entrypoint:** `src/vuln_analysis/functions/cve_agent.py` + +Before any checklist question is investigated, the orchestrator prepares the execution environment: + +| Step | Description | +|---|---| +| **Load workflow state** | Reads `checklist_plans` (CVE ID → list of questions) from `ExploitIqEngineState`. | +| **Instantiate sub-agents** | Iterates the agent registry (`reachability`, `code_understanding`), loads each agent's tool set, and compiles its LangGraph state machine. | +| **Fan out across CVEs** | Uses `asyncio.gather` to process all CVEs in parallel. Within each CVE, questions also run concurrently (optionally limited by a semaphore). | +| **Build intelligence** | For each CVE, extracts critical context, candidate packages, and vulnerable functions from pre-computed CVE intel. Go ecosystems get additional candidate enrichment; vulnerable functions may be enriched from patch data. | +| **Synthetic reachability injection** | If no question was routed to reachability but candidate packages exist, a synthetic reachability question is auto-injected to ensure at least one call-chain check runs. | + +### 2. Dispatcher (Question Routing) + +**File:** `src/vuln_analysis/functions/dispatcher.py` + +For each checklist question, a routing LLM classifies the question into one of two sub-agent types using structured output (`QuestionRouting`): + +| Route | Sub-Agent | When to use | +|---|---|---| +| `reachability` | **Transitive Search / Reachability Agent** | Call chains, code paths, whether vulnerable code is called/used/reachable, whether untrusted data can reach a function. | +| `code_understanding` | **Code Understanding Agent** (includes config/env tasks) | Configuration, version, presence, application-level settings, input validation, and general behavioral questions that do not require call-chain tracing. | + +The dispatcher prompt includes 14 labeled examples. On routing failure, the system falls back to `reachability`. + +### 3. Sub-Agent Execution Loop + +Once routed, the selected sub-agent receives the checklist question and runs its investigation loop. This is the core of the system. + +Each sub-agent inherits a shared five-node LangGraph state machine from `BaseGraphAgent` (`src/vuln_analysis/functions/base_graph_agent.py`). Sub-agents extend the template through abstract methods and optional hooks while sharing the same control flow. + +#### Shared State Machine Nodes + +| Node | LLM? | Role | +|---|---|---| +| **PreProcess** | No | Select target package, build ecosystem-specific system prompt, initialize rules tracker. | +| **Thought** | Yes (structured) | Decide next action (`act` or `finish`) via Pydantic `Thought` schema. | +| **Tool** | No | Execute the selected reachability or code-understanding tool. | +| **Observation** | Yes (structured) | Enforce rules, compress tool output via comprehension + memory update LLMs, prune tokens. | +| **ForcedFinish** | Yes | Safety net when `max_iterations` is reached; generates a final answer from accumulated memory. | + +#### Node Details + +**`pre_process_node`** - Each sub-agent implements this to initialize its investigation: +- Selects the target package from candidates (LLM filter, Go stdlib matching, or image/repo matching) +- Builds the ecosystem-specific system prompt with tool descriptions and instructions +- Configures the rules tracker with target package, allowed tools, and vulnerable functions + +**`thought_node`** - The reasoning engine. At each step: +1. Assembles messages: system prompt + conversation history + observation context (KNOWLEDGE block) +2. Prunes old messages to fit within the token budget +3. Invokes the LLM with `Thought` structured output +4. On `mode="act"`: builds tool call arguments and creates an `AIMessage` with `tool_calls` +5. On `mode="finish"`: checks `check_finish_allowed()` hook - if blocked, injects error message and loops back +6. On missing actions: forces a finish with "Insufficient evidence" + +**`tool_node`** - Standard LangGraph `ToolNode` that executes the selected tool and returns the raw `ToolMessage`. + +**`observation_node`** - The semantic compression engine. After each tool call: +1. **Rule enforcement** - calls `rules_tracker.check_thought_behavior()`; on violation, returns an error message instead of processing tool output +2. **Comprehension** - invokes the comprehension LLM to extract 3–5 key technical facts from raw tool output +3. **CVE ID sanitization** - replaces hallucinated CVE IDs with "the investigated vulnerability" +4. **Memory update** - invokes the memory update LLM to merge new findings into cumulative memory +5. **Token pruning** - removes old messages when estimated tokens exceed the budget +6. **Post-observation hook** - calls `post_observation()` for agent-specific processing + +**`forced_finish_node`** — Triggered when `step >= max_steps`. Generates a final answer under token constraints. The reachability agent overrides this to inject a warning when Call Chain Analyzer was never called. + +**Routing (`should_continue`):** +- `step >= max_steps` → `forced_finish_node` +- `thought is None` → `thought_node` (re-entry after rule violation) +- `thought.mode == "finish"` → `END` +- `thought.mode == "act"` → `tool_node` + +**Internal loop:** `Thought → Tool → Observation → Thought` repeats until the agent finishes, hits the iteration limit, or a rule violation forces a re-think. + +#### Sub-Agent: Reachability + +**File:** `src/vuln_analysis/functions/reachability_agent.py` + +Traces call chains to determine if vulnerable code is **reachable from application code** — the critical distinction between code being present in a container and code being exploitable on an actual execution path. + +**Tools:** + +| Tool | Purpose | +|---|---| +| Function Locator | Validates package/function names with fuzzy matching | +| Call Chain Analyzer | Determines if a function is reachable from app code | +| Function Library Version Finder | Checks installed library version (Java only) | +| Function Caller Finder | Finds functions calling stdlib methods (Go only) | +| Code Keyword Search | Exact keyword matching in source code | +| CVE Web Search | External CVE/vulnerability information lookup | +| Docs Semantic Search | Vector embeddings search in documentation | + +**Investigation protocol:** +1. IDENTIFY the vulnerable component/function from the CVE description +2. SEARCH for its presence using Code Keyword Search +3. TRACE reachability using Call Chain Analyzer (Function Locator first; Function Caller Finder before CCA for Go) +4. ASSESS exploitability only after completing reachability checks + +Key behavioral rules: Code Keyword Search proves **presence**, not reachability. Only Call Chain Analyzer can confirm reachability. The agent must distinguish PRESENT vs. REACHABLE vs. EXPLOITABLE. For non-reachability questions routed to this agent, CCA and FCF tools are removed and a simpler prompt is used. + +#### Sub-Agent: Code Understanding + +**File:** `src/vuln_analysis/functions/code_understanding_agent.py` + +Investigates configuration, version, presence, and application-level behavior **without tracing call chains**. + +**Tools:** + +| Tool | Purpose | +|---|---| +| Docs Semantic Search | Semantic search in documentation | +| Code Keyword Search | Exact keyword matching in source code | +| Configuration Scanner | Scans YAML, XML, properties, build files | +| Import Usage Analyzer | Finds all imports/usage of a package | + +Key rules: a component in dependencies does not mean the application uses it. Configuration in framework dependencies and imports in intermediate libraries are relevant evidence. The agent uses minimal grounding context during comprehension to prevent CVE ID hallucination from parametric model knowledge. + +#### Rule System + +Rules are checked in `observation_node` before tool output is processed. On violation, the agent loops back to `thought_node` with an error message. + +**Base rules (all sub-agents):** +- Duplicate call — same tool + same input as a previous call +- Dotted query retry — Code Keyword Search with dots returned empty after a previous dotted-empty attempt +- Allowed tools — tool not in the agent's permitted set + +**Reachability-specific rules:** +- Target package first — FL/CCA/FCF called with wrong package before using target package +- Target functions first — CCA/FCF called with non-priority function when GHSA functions are unchecked +- CCA mandatory — agent tries to finish without ever calling CCA +- Java FLVF — CCA found reachability but Function Library Version Finder not called + +#### Semantic Compression Pipeline + +Each tool call produces raw output that may be thousands of tokens. The observation node compresses this through a two-stage LLM pipeline: + +1. **Comprehension LLM** — extracts 3–5 key technical facts from raw tool output into structured `CodeFindings` +2. **Memory update LLM** — merges new findings into cumulative `Observation` memory, deduplicating and recording tool call history + +Token budget is managed at multiple levels: comprehension prompt truncation, observation context capping (1200 tokens), message pruning in thought/observation nodes, and a fallback to truncated raw excerpts on comprehension failures. + +#### Memory System + +`InternalMemory` provides deterministic, token-budgeted storage for agent observations. Each `MemoryItem` has a category (e.g. `reachable`, `not_reachable`, `vuln_found`, `tool_record`), priority (1–3), and deduplication key. When memory exceeds its token budget, lowest-priority items are evicted first (FIFO within equal priority), ensuring high-priority findings like reachability results survive. + +### 4. Checklist Iteration + +After a sub-agent completes one checklist question (reaches `END` or `forced_finish_node`), the orchestrator checks whether more questions remain for the current CVE. If yes, the dispatcher routes the next question. This outer loop continues until every question across every CVE has been answered. + +### 5. Post-Process + +**Function:** `_postprocess_results()` in `cve_agent.py` + +Once all sub-agent runs complete, the orchestrator normalizes raw graph outputs into a uniform structure: + +| Field | Description | +|---|---| +| `input` | The checklist question text | +| `output` | The sub-agent's final answer | +| `cca_results` | Call Chain Analyzer verdicts (reachability agent) | +| `package_validated` | Whether Function Locator validated the target package | +| `is_reachability` | Whether the question was classified as a reachability investigation | + +Exceptions during agent execution can optionally be replaced with a configured placeholder message (`replace_exceptions`). The final `checklist_results` dict (CVE ID → list of answer dicts) is written back to workflow state for downstream nodes (summarization, justification, CVSS generation). + +--- + +## Process Diagrams + +The investigation runs at two levels. The **checklist loop** routes each question to the right specialist and repeats until every item is answered. Inside each routed question, the **sub-agent state machine** autonomously investigates in a repeating cycle: decide what to search → run the search → summarize what was found. + +### Checklist Investigation Flow + +For each checklist question, the dispatcher analyzes the question and routes it to a specialized sub-agent. The sub-agent investigates until the question is answered, then the system moves to the next checklist item. + +```mermaid +%%{init: {'theme': 'base', 'themeVariables': {'primaryTextColor': '#000000', 'secondaryTextColor': '#000000', 'tertiaryTextColor': '#000000', 'lineColor': '#333333'}}}%% +flowchart TD + StartChecklist([Start — next checklist question]) + + StartChecklist --> Dispatcher{"Dispatcher
Analyzes the question"} + Dispatcher -->|Route to specialist| TaskType{"What type of task is this?"} + + TaskType -->|"Reachability - trace code paths"| ReachabilityAgent["Reachability Sub-Agent
Traces call chains and validates
function reachability from application code"] + TaskType -->|"Code understanding"| CodeUnderstandingAgent["Code Understanding Sub-Agent
Investigates behavior, imports, and presence
without tracing call chains"] + + ReachabilityAgent --> RunInvestigation + CodeUnderstandingAgent --> RunInvestigation + + RunInvestigation["Run investigation in sub-agent
Executes the investigation loop below
until the question is answered"] + + RunInvestigation --> TaskComplete{"Is the question answered?"} + TaskComplete -->|No — continue investigating| RunInvestigation + TaskComplete -->|Yes| MoreItems{"More checklist items?"} + MoreItems -->|Yes| StartChecklist + MoreItems -->|No| EndFlow([All checklist items complete]) + + style StartChecklist fill:#ececec,stroke:#2d2d2d,color:#000000 + style Dispatcher fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style TaskType fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style TaskComplete fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style MoreItems fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style ReachabilityAgent fill:#c8e0c8,stroke:#4d8f4d,color:#000000 + style CodeUnderstandingAgent fill:#c8e0c8,stroke:#4d8f4d,color:#000000 + style RunInvestigation fill:#c8e0c8,stroke:#4d8f4d,color:#000000 + style EndFlow fill:#ececec,stroke:#2d2d2d,color:#000000 +``` + +### Sub-Agent State Machine + +Once a question is routed, the selected sub-agent runs an investigation loop. **Blue** steps use an AI model. **Green** steps run deterministic code. + +```mermaid +%%{init: {'theme': 'base', 'themeVariables': {'primaryTextColor': '#000000', 'secondaryTextColor': '#000000', 'tertiaryTextColor': '#000000', 'lineColor': '#333333'}}}%% +flowchart TD + StartInvestigation([Start investigation]) + + StartInvestigation --> PrepareInvestigation + + PrepareInvestigation["Prepare the investigation
Gather CVE context and candidate packages
Choose the right approach for the codebase
Set up the question and investigation rules"] + PrepareInvestigation --> DecideNextStep + + DecideNextStep["Decide what to do next
The AI reviews what is known so far
and chooses the next search or the final answer"] + + DecideNextStep -.->|Need more evidence| RunCodeSearch + RunCodeSearch["Run a code analysis search
Execute the chosen search tool
such as tracing call chains,
searching source code, or reading config files"] + RunCodeSearch --> SummarizeFindings + + SummarizeFindings["Summarize what was found
The AI reads the search results,
checks they follow the investigation rules,
and saves the important findings to memory"] + + SummarizeFindings --> DecideNextStep + + DecideNextStep -.->|Enough evidence to answer| InvestigationComplete + DecideNextStep -.->|Investigation step limit reached| WriteFinalAnswer + + WriteFinalAnswer["Write a final answer from what was found
If the investigation ran out of steps,
the AI produces the best possible answer
from everything collected so far"] + + WriteFinalAnswer --> InvestigationComplete + InvestigationComplete([Investigation complete]) + + style DecideNextStep fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style SummarizeFindings fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style WriteFinalAnswer fill:#c5daf0,stroke:#3d6ea8,color:#000000 + style PrepareInvestigation fill:#c8e0c8,stroke:#4d8f4d,color:#000000 + style RunCodeSearch fill:#c8e0c8,stroke:#4d8f4d,color:#000000 + style StartInvestigation fill:#ececec,stroke:#2d2d2d,color:#000000 + style InvestigationComplete fill:#ececec,stroke:#2d2d2d,color:#000000 +``` + +--- + +## CVE Agent Loop Architecture (Detailed Reference) + +The CVE Agent Loop is the core investigation stage of the exploit-iq-agent pipeline. It receives a set of checklist questions per CVE and dispatches each question to a specialized sub-agent that autonomously investigates the answer using code analysis tools, structured reasoning, and behavioral rules. + +### Table of Contents + +- [High-Level Overview](#high-level-overview) +- [Orchestration Layer](#orchestration-layer) + - [Entrypoint](#entrypoint) + - [Question Dispatcher](#question-dispatcher) + - [Synthetic Reachability Injection](#synthetic-reachability-injection) +- [Agent Registry](#agent-registry) +- [Shared State Machine (Template Algorithm)](#shared-state-machine-template-algorithm) + - [Agent State](#agent-state) + - [Graph Topology](#graph-topology) + - [Node Descriptions](#node-descriptions) + - [Template Method & Hook Points](#template-method--hook-points) +- [Sub-Agent: Reachability](#sub-agent-reachability) + - [Purpose & Tool Set](#reachability-purpose--tool-set) + - [System Prompt](#reachability-system-prompt) + - [Ecosystem-Specific Instructions](#ecosystem-specific-instructions) + - [Rules Tracker](#reachability-rules-tracker) + - [Behavioral Overrides](#reachability-behavioral-overrides) +- [Sub-Agent: Code Understanding](#sub-agent-code-understanding) + - [Purpose & Tool Set](#code-understanding-purpose--tool-set) + - [System Prompt](#code-understanding-system-prompt) + - [Rules Tracker](#code-understanding-rules-tracker) + - [Anti-Hallucination Design](#anti-hallucination-design) +- [Rule System](#rule-system) + - [Base Rules](#base-rules) + - [Reachability-Specific Rules](#reachability-specific-rules) + - [Enforcement Mechanism](#enforcement-mechanism) +- [Tool Set](#tool-set) + - [Tool Availability](#tool-availability) +- [Semantic Compression Pipeline](#semantic-compression-pipeline) + - [Comprehension LLM](#comprehension-llm) + - [Memory Update LLM](#memory-update-llm) + - [Token Budget Management](#token-budget-management) +- [Structured Output Schemas](#structured-output-schemas) +- [Memory System](#memory-system) + +--- + +### High-Level Overview + +``` +┌──────────────────────────────────────────────────────────────────────┐ +│ cve_agent (orchestrator) │ +│ │ +│ For each CVE: │ +│ ┌────────────────┐ ┌─────────────┐ ┌───────────────────────┐ │ +│ │ Build Intel │───▶│ Dispatcher │───▶│ Execute Sub-Agents │ │ +│ │ (context, pkgs, │ │ (LLM-based │ │ (parallel, semaphore │ │ +│ │ vuln functions)│ │ routing) │ │ controlled) │ │ +│ └────────────────┘ └─────────────┘ └───────────────────────┘ │ +│ │ │ │ +│ ┌─────────┴─────────┐ │ │ +│ ▼ ▼ ▼ │ +│ ┌──────────────┐ ┌─────────────────┐ ┌──────────┐ │ +│ │ Reachability │ │Code Understanding│ │ Post- │ │ +│ │ Agent │ │Agent │ │ Process │ │ +│ └──────────────┘ └─────────────────┘ └──────────┘ │ +└──────────────────────────────────────────────────────────────────────┘ +``` + +The system follows a **fan-out/fan-in** pattern: all CVE investigation groups run concurrently via `asyncio.gather`, and within each group, individual checklist questions run concurrently (with optional semaphore-based concurrency control). + +--- + +### Orchestration Layer + +#### Entrypoint + +**File:** `src/vuln_analysis/functions/cve_agent.py` + +The `cve_agent` function is the pipeline stage entrypoint registered via `@register_function`. It: + +1. **Instantiates sub-agents** by iterating over all registered agent types, loading their tools, and compiling their LangGraph state machines. +2. **Fans out across CVEs** using `asyncio.gather` — each CVE's checklist questions are processed by `_process_steps()` concurrently. +3. **Collects and post-processes results** into `state.checklist_results`. + +```python +## Simplified flow +agents = {} +for agent_type in get_all_agent_types(): + agent_cls = get_agent_class(agent_type) + tools = agent_cls.get_tools(builder, config, state) + agent_instance = agent_cls(tools=tools, llm=llm, config=config) + agents[agent_type] = await agent_instance.build_graph() + +results = await asyncio.gather( + *(_process_steps(agents, routing_llm, steps, semaphore, ..., vuln_id=vid) + for vid, steps in checklist_plans.items()), + return_exceptions=True, +) +``` + +**Configuration** is defined by `CVEAgentExecutorToolConfig`: + +| Parameter | Default | Description | +|---|---|---| +| `llm_name` | — | LLM model identifier | +| `max_concurrency` | None | Semaphore limit for parallel agent runs | +| `max_iterations` | 10 | Maximum thought-act-observe cycles per question | +| `context_window_token_limit` | 5000 | Token budget for observation context pruning | +| `uber_jar_file_threshold` | 600 | JAR file count above which uber-JAR handling activates | +| `cve_web_search_enabled` | True | Feature flag for CVE Web Search tool | +| `transitive_search_tool_enabled` | True | Feature flag for Function Locator / CCA / FCF | + +#### Question Dispatcher + +**File:** `src/vuln_analysis/functions/dispatcher.py` + +Before execution, each checklist question is classified by an LLM with structured output into one of two agent types: + +``` +┌───────────────────────────┐ +│ Checklist Question │ +└─────────────┬─────────────┘ + │ + ▼ +┌───────────────────────────┐ +│ Routing LLM (structured │ +│ output: QuestionRouting) │ +└─────────┬─────────┬───────┘ + │ │ + ▼ ▼ + "reachability" "code_understanding" +``` + +**Routing criteria:** + +- **Reachability:** call chains, code path tracing, whether vulnerable code is called/used/reachable, whether untrusted data can reach a function. +- **Code Understanding:** configuration, version, presence, application-level settings, few code understanding capabilities (like input validation for example) + +The dispatcher prompt includes 14 labeled examples to anchor the LLM's classification. If routing fails (exception), it falls back to `"reachability"`. + +#### Synthetic Reachability Injection + +If no question in a CVE's checklist was routed to the reachability agent but candidate packages exist, the orchestrator auto-injects a synthetic reachability question: + +> "Is the code from the vulnerable package '{name}' actually imported, called, or reachable from the application's own source code?" + +This ensures every CVE investigation includes at least one reachability check when there are packages to investigate. + +--- + +### Agent Registry + +**File:** `src/vuln_analysis/functions/agent_registry.py` + +A lightweight decorator-based registry that maps string type names to `BaseGraphAgent` subclasses: + +```python +@register_agent("reachability") +class ReachabilityAgent(BaseGraphAgent): ... + +@register_agent("code_understanding") +class CodeUnderstandingAgent(BaseGraphAgent): ... +``` + +The registry provides two functions: +- `get_agent_class(agent_type)` — retrieve a class by name +- `get_all_agent_types()` — list all registered type strings + +Agent modules are imported in `cve_agent.py` to trigger the `@register_agent` decorators at module load time. + +--- + +### Shared State Machine (Template Algorithm) + +**File:** `src/vuln_analysis/functions/base_graph_agent.py` + +`BaseGraphAgent` is an abstract base class implementing the **Template Method pattern** over a LangGraph `StateGraph`. It defines the shared graph topology, node implementations, and control flow. Sub-agents inherit the entire execution loop and extend it through abstract methods and optional hooks. + +#### Agent State + +**File:** `src/vuln_analysis/functions/react_internals.py` + +`AgentState` extends LangGraph's `MessagesState` with investigation-specific fields: + +```python +class AgentState(MessagesState): + input: str # The checklist question + step: int # Current iteration counter + max_steps: int # Maximum allowed iterations + thought: Thought | None # LLM's structured reasoning output + observation: Observation | None # Accumulated findings & memory + output: str # Final answer text + ecosystem: str | None # "go", "python", "java", etc. + runtime_prompt: str | None # Agent-specific system prompt + app_package: str | None # Target package under investigation + is_reachability: str # "yes" or "no" + rules_tracker: BaseRulesTracker # Behavioral rule enforcer + critical_context: list[str] # CVE metadata & advisory context + cca_results: list[bool] # Call Chain Analyzer verdicts + package_validated: bool | None # Function Locator validation + precomputed_intel: tuple | None # Pre-extracted CVE intelligence +``` + +#### Graph Topology + +``` + ┌──────────────────────────────┐ + │ START │ + └──────────────┬───────────────┘ + │ + ▼ + ┌──────────────────────────────┐ + │ pre_process_node │ + │ (ABSTRACT — subclass impl) │ + │ • Select target package │ + │ • Build system prompt │ + │ • Initialize rules tracker │ + └──────────────┬───────────────┘ + │ + ┌─────────────▼───────────────┐ + ┌───▶│ thought_node │◀──────────────┐ + │ │ • LLM decides: act/finish │ │ + │ │ • Structured output: Thought │ │ + │ └──────────┬───────────────────┘ │ + │ │ │ + │ ┌────────┴─────────┬──────────────┐ │ + │ │ │ │ │ + │ ▼ ▼ ▼ │ + │ mode="finish" mode="act" step >= max │ + │ │ │ │ │ + │ ▼ ▼ ▼ │ + │ ┌──────┐ ┌──────────────┐ ┌────────────┐ │ + │ │ END │ │ tool_node │ │forced_finish│ │ + │ └──────┘ │ (ToolNode) │ │ _node │ │ + │ └──────┬───────┘ └─────┬──────┘ │ + │ │ │ │ + │ ▼ ▼ │ + │ ┌─────────────────┐ ┌──────┐ │ + │ │observation_node │ │ END │ │ + │ │ • Rule check │ └──────┘ │ + │ │ • Comprehension │ │ + │ │ • Memory update │ │ + │ │ • Token pruning │ │ + └──────────────┤ • Post-obs hook │ │ + └──────────────────┘ │ + │ │ + └───────────────────────────┘ +``` + +#### Node Descriptions + +##### `pre_process_node` (Abstract) + +Each sub-agent implements this to initialize its investigation: +- Selects the target package from candidates (using LLM filter, Go stdlib matching, or image/repo matching) +- Builds the ecosystem-specific system prompt with tool descriptions and instructions +- Configures the rules tracker with target package, allowed tools, and vulnerable functions +- Returns the initialized state + +##### `thought_node` (Shared) + +The reasoning engine. At each step: +1. Assembles messages: system prompt + conversation history + observation context (KNOWLEDGE block) +2. Prunes old messages to fit within the token budget +3. Invokes the LLM with `Thought` structured output +4. On `mode="act"`: builds tool call arguments and creates an `AIMessage` with `tool_calls` +5. On `mode="finish"`: checks `check_finish_allowed()` hook — if blocked, injects error message and loops back +6. On missing actions: forces a finish with "Insufficient evidence" + +##### `tool_node` (LangGraph ToolNode) + +Standard LangGraph `ToolNode` — executes the selected tool and returns the raw `ToolMessage`. + +##### `observation_node` (Shared) + +The semantic compression engine. After each tool call: +1. **Rule enforcement:** calls `rules_tracker.check_thought_behavior()` — if violated, returns error message to the agent instead of processing +2. **Comprehension:** invokes the comprehension LLM to extract 3-5 key technical facts from raw tool output +3. **CVE ID sanitization:** replaces hallucinated CVE IDs with "the investigated vulnerability" +4. **Memory update:** invokes the memory update LLM to merge new findings into cumulative memory +5. **Token pruning:** removes old messages when estimated tokens exceed the budget +6. **Post-observation hook:** calls `post_observation()` for agent-specific processing + +##### `forced_finish_node` (Overridable) + +Triggered when `step >= max_steps`. Generates a final answer under token constraints. The reachability agent overrides this to inject a "no CCA warning" when Call Chain Analyzer was never called. + +##### `should_continue` (Routing) + +Determines the next node: +- `step >= max_steps` → `forced_finish_node` +- `thought is None` → `thought_node` (re-entry after rule violation) +- `thought.mode == "finish"` → `END` +- `thought.mode == "act"` → `tool_node` + +#### Template Method & Hook Points + +`BaseGraphAgent` defines three **abstract methods** (must be implemented) and five **optional hooks** (may be overridden): + +``` +BaseGraphAgent (Template) +│ +├── Abstract (MUST implement): +│ ├── pre_process_node(state) → AgentState +│ ├── get_tools(builder, config, state) → list +│ └── create_rules_tracker() → BaseRulesTracker +│ +├── Optional Hooks (MAY override): +│ ├── build_comprehension_context(state) → str +│ ├── sanitize_findings(findings, state) → list[str] +│ ├── post_observation(state, tool_used, ...) → dict +│ ├── should_truncate_tool_output(state, tool_used) → bool +│ └── check_finish_allowed(state) → (bool, str) +│ +└── Shared Concrete (inherited as-is): + ├── thought_node(state) → AgentState + ├── observation_node(state) → AgentState + ├── forced_finish_node(state) → AgentState (overridable) + ├── should_continue(state) → str + └── build_graph() → CompiledGraph +``` + +--- + +### Sub-Agent: Reachability + +**File:** `src/vuln_analysis/functions/reachability_agent.py` +**Registry name:** `"reachability"` + +#### Reachability Purpose & Tool Set + +Traces call chains to determine if vulnerable code is **reachable from application code**. This is the critical distinction between code being present in a container and code being exploitable ( Only if the code present in the actual execution path). + +**Tools (8):** + +| Tool | Purpose | +|---|---| +| Function Locator | Validates package/function names with fuzzy matching | +| Call Chain Analyzer | Determines if a function is reachable from app code | +| Function Library Version Finder | Checks installed library version (Java only) | +| Function Caller Finder | Finds functions calling stdlib methods (Go only) | +| Code Keyword Search | Exact keyword matching in source code | +| CVE Web Search | External CVE/vulnerability information lookup | +| Docs Semantic Search | Vector embeddings search in documentation | + +#### Reachability System Prompt + +The reachability system prompt enforces a **4-step mandatory investigation protocol:** + +``` +1. IDENTIFY the vulnerable component/function from the CVE description +2. SEARCH for its presence using Code Keyword Search +3. TRACE reachability using Call Chain Analyzer + - Use Function Locator first to verify names + - For Go: use Function Caller Finder before CCA +4. ASSESS exploitability only after completing reachability checks +``` + +Key rules embedded in the system prompt: +- Code Keyword Search proves **presence**, NOT reachability +- Function Locator validates packages **names**, NOT reachability +- Only Call Chain Analyzer tool can confirm reachability +- Positive CCA (Call Chain Analyzer) (True) → may conclude exploitable +- Negative CCA (False) → record result, may conclude not exploitable +- Must distinguish PRESENT vs. REACHABLE vs. EXPLOITABLE + +The prompt also includes the thought instructions with numbered rules and few-shot examples showing the expected investigation flow (search → locate → trace). + +#### Ecosystem-Specific Instructions + +The reachability agent adapts its instructions based on ecosystem: + +| Ecosystem | Variant | Key Difference | +|---|---|---| +| General | `REACHABILITY_AGENT_THOUGHT_INSTRUCTIONS` | 7 rules + 3 examples | +| Go | `REACHABILITY_AGENT_THOUGHT_INSTRUCTIONS_GO` | Rule 8: FCF before CCA, 4 examples | +| Python | `REACHABILITY_AGENT_THOUGHT_INSTRUCTIONS_PYTHON` | Rule 8: class names for CCA (not methods) | +| Non-reach | `REACHABILITY_AGENT_NON_REACH_THOUGHT_INSTRUCTIONS` | No CCA/FCF, focus on presence/config | + +For **non-reachability questions** routed to this agent, CCA and FCF tools are removed and a different system prompt is used that focuses on presence/version/configuration without call chain tracing. + +#### Reachability Rules Tracker + +`ReachabilityRulesTracker` enforces 6 rules (see [Rule System](#rule-system)). + +#### Reachability Behavioral Overrides + +**`pre_process_node`:** Classifies the question as reachability or non-reachability using an LLM. For non-reachability, removes CCA/FCF tools and switches to a simpler prompt. + +**`forced_finish_node`:** When the agent exhausts all iterations without ever calling CCA on a reachability question, injects a strong warning: + +> "Call Chain Analyzer was NEVER called during this investigation. Without Call Chain Analyzer verification, there is NO evidence that the vulnerable function is reachable. Library presence alone does NOT constitute exploitability." + +This prevents the LLM from hallucinating exploitability based on library presence alone. + +**`post_observation`:** Extracts CCA results (True/False) and tracks package validation status for uber-JAR handling. + +**`check_finish_allowed`:** Blocks premature finish if CCA was never called, or if (Java only) CCA found reachability but version was not verified via FLVF. + +**`should_truncate_tool_output`:** Enables Java-specific truncation of verbose tool output. + +--- + +### Sub-Agent: Code Understanding + +**File:** `src/vuln_analysis/functions/code_understanding_agent.py` +**Internals:** `src/vuln_analysis/functions/code_understanding_internals.py` +**Registry name:** `"code_understanding"` + +#### Code Understanding Purpose & Tool Set + +Investigates configuration, version, presence, and application-level behavior **without tracing call chains**. +It also has some modest code understanding capabilities, mainly from the LLM being used. +There are future plans to extend CU (Code understanding) sub-agent' toolset to include static code analysis tools +to further improve the code understanding capabilities. + +**Tools (4):** + +| Tool | Purpose | +|---|---| +| Docs Semantic Search | Semantic search in documentation | +| Code Keyword Search | Exact keyword matching in source code | +| Configuration Scanner | Scans YAML, XML, properties, build files | +| Import Usage Analyzer | Finds all imports/usage of a package | + +#### Code Understanding System Prompt + +The code understanding prompt focuses on evidence-based investigation: + +``` +"This is NOT a reachability question -- do NOT trace call chains." +``` + +Key rules: +- A component being present in dependencies does NOT mean the application uses it +- Configuration in framework dependencies IS relevant evidence +- Import in an intermediate library IS relevant +- Synthesize findings across documentation, code, configuration, and imports + +The thought instructions include 9 rules and 3 examples, emphasizing the use of Configuration Scanner first, followed by Code Keyword Search and Import Usage Analyzer. + +#### Code Understanding Rules Tracker + +`CodeUnderstandingRulesTracker` enforces 3 rules from the base class (see [Rule System](#rule-system)). + +#### Anti-Hallucination Design + +The code understanding agent overrides `build_comprehension_context()` to return **minimal grounding context**: + +```python +def build_comprehension_context(self, state): + return ( + f"Investigating {vuln_id} in package {package}.\n" + "Only extract facts explicitly stated in the tool output. " + "Do not add CVE IDs, vulnerability names, or advisory details " + "from your own knowledge." + ) +``` + +Feeding the full GHSA/NVD context activates the model's parametric knowledge and causes it to inject CVE IDs from training data. The minimal context prevents this hallucination pathway. + +--- + +### Rule System + +The rule system provides deterministic behavioral guardrails that prevent the LLM from making common investigation mistakes. Rules are checked **before** tool output is processed in `observation_node`. + +#### Base Rules + +Shared by all sub-agents via `BaseRulesTracker`: + +| Rule | ID | Trigger | Error Message | +|---|---|---|---| +| **Duplicate Call** | 0 | Same tool + same input as a previous call | "You already called {tool} with this exact input" | +| **Dotted Query Retry** | 5 | Code Keyword Search with dots returns empty after a previous dotted-empty | "Your query contains dots and returned no results. Retry with final component" | +| **Allowed Tools** | — | Tool not in the agent's allowed set | "You MUST use the allowed tools {list}" | + +#### Reachability-Specific Rules + +Additional rules enforced by `ReachabilityRulesTracker`: + +| Rule | ID | Trigger | Error Message | +|---|---|---|---| +| **Target Package First** | 6 | FL/CCA/FCF called with wrong package before using target package | "You MUST use the target package name {pkg} first" | +| **Target Functions First** | 7 | CCA/FCF called with non-priority function when GHSA functions are unchecked | "You MUST investigate {functions} first" | +| **CCA Mandatory** | 8 | Agent tries to finish without ever calling CCA | "You MUST use Function Locator and Call Chain Analyzer before concluding" | +| **Java FLVF** | 9 | (Java) CCA found reachability but FLVF not called | "You MUST call Function Library Version Finder to check version" | + +#### Enforcement Mechanism + +``` +thought_node ──▶ tool_node ──▶ observation_node + │ + ┌──────┴───────┐ + │ rules_tracker │ + │ .check_ │ + │ thought_ │ + │ behavior() │ + └──────┬───────┘ + │ + ┌─────────┴─────────┐ + │ │ + violation no violation + │ │ + ▼ ▼ + Return error msg Process tool output + as HumanMessage (comprehension, + → loops back to memory, etc.) + thought_node +``` + +When a rule is violated, the observation node returns the error message as a `HumanMessage` without processing the tool output. The agent loops back to `thought_node` where it must correct its behavior. + +--- + +### Tool Set + +**File:** `src/vuln_analysis/tools/tool_names.py` + +| Tool Name | Ecosystems | Category | Description | +|---|---|---|---| +| **Code Semantic Search** | All | Search | Searches source code using semantic vector embeddings | +| **Docs Semantic Search** | All | Search | Searches documentation using semantic vector embeddings | +| **Code Keyword Search** | All | Search | Exact keyword matching in source code | +| **Function Locator** | All | Code Path | Validates package/function names with fuzzy matching | +| **Call Chain Analyzer** | All | Code Path | Determines if a function is reachable from application code | +| **Function Caller Finder** | Go | Code Path | Finds functions calling specific stdlib methods | +| **Function Library Version Finder** | Java | Version | Checks installed library version | +| **CVE Web Search** | All | External | Web search for CVE/vulnerability information | +| **Configuration Scanner** | All | Config | Scans config files for security-relevant patterns | +| **Import Usage Analyzer** | All | Usage | Finds all imports/usage of a package across sources | +| **Container Analysis Data** | All | Pre-computed | Pre-analyzed data from earlier pipeline stages | +| **Source Grep** | All | Search | Native Unix grep on source code | + +#### Tool Availability + +Tools are conditionally available based on runtime state and configuration: + +```python +_TOOL_AVAILABILITY = { + "Code Semantic Search": lambda c, s: s.code_vdb_path is not None, + "Docs Semantic Search": lambda c, s: s.doc_vdb_path is not None, + "Code Keyword Search": lambda c, s: s.code_index_path is not None, + "Import Usage Analyzer": lambda c, s: s.code_index_path is not None, + "CVE Web Search": lambda c, s: c.cve_web_search_enabled, + "Call Chain Analyzer": lambda c, s: c.transitive_search_tool_enabled and s.code_index_path is not None, + "Function Caller Finder": lambda c, s: c.transitive_search_tool_enabled and s.code_index_path is not None, + "Function Locator": lambda c, s: c.transitive_search_tool_enabled and s.code_index_path is not None, +} +``` + +Tools not in this map (e.g., Configuration Scanner) are always available. Each sub-agent further filters to its own tool subset. + +--- + +### Semantic Compression Pipeline + +Each tool call produces raw output that may be thousands of tokens. The observation node compresses this through a two-stage LLM pipeline before storing it in memory. + +``` +Raw Tool Output + │ + ▼ +┌─────────────────────────────────────────────┐ +│ Comprehension LLM │ +│ Input: goal, package, context, tool output │ +│ Output: CodeFindings │ +│ • findings: 3-5 key technical facts │ +│ • tool_outcome: "CALLED: X with Y → Z" │ +└──────────────────┬──────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────┐ +│ Memory Update LLM │ +│ Input: goal, package, previous memory, │ +│ new findings, tool outcome │ +│ Output: Observation │ +│ • results: current step's findings │ +│ • memory: cumulative deduplicated facts │ +└──────────────────┬──────────────────────────┘ + │ + ▼ + Updated AgentState +``` + +#### Comprehension LLM + +**Prompt:** `COMPREHENSION_PROMPT` + +Extracts structured facts from raw tool output with code comprehension rules: +- Read actual code snippets, don't just note presence +- Determine what code DOES (functional behavior), not just where it is +- Distinguish same-named entities in different packages +- Check technology fit (language/framework mismatches) +- Ground findings to vendor mitigations when available + +Special handling for Call Chain Analyzer results: +- Positive (reachable): first finding MUST be `"REACHABLE via [package] - sufficient evidence"` +- Negative (not reachable): first finding MUST be `"NOT reachable via [package]"` + +#### Memory Update LLM + +**Prompt:** `MEMORY_UPDATE_PROMPT` + +Merges new findings into cumulative investigation memory: +- Start from previous memory, append new facts +- Record absence if nothing was found +- Add tool call record verbatim (so future steps know what was tried) +- Preserve "Main application" vs. "Application library dependencies" distinction +- Check findings against known mitigations from the vulnerability context + +#### Token Budget Management + +The system manages tokens at multiple levels: + +1. **Comprehension prompt:** truncates tool output if the prompt exceeds `context_window_token_limit - 2000` +2. **Observation context:** the KNOWLEDGE block shown to thought_node is capped at `_OBS_CONTEXT_MAX_TOKENS` (1200 tokens) +3. **Message pruning:** old conversation messages are removed (preserving system prompt and tail) when estimated tokens exceed the configured limit +4. **Comprehension fallback:** if the comprehension LLM hits a `LengthFinishReasonError`, returns a truncated raw excerpt instead of crashing + +--- + +### Structured Output Schemas + +All LLM interactions use Pydantic models as structured output schemas, preventing free-form hallucination: + +``` +┌──────────────────────────────────────────────────────┐ +│ Thought │ +│ thought: str (max 3000 chars) │ +│ mode: "act" | "finish" │ +│ actions: ToolCall | None (when mode="act") │ +│ final_answer: str | None (when mode="finish") │ +├──────────────────────────────────────────────────────┤ +│ ToolCall │ +│ tool: str (from AVAILABLE_TOOLS) │ +│ package_name: str | None │ +│ function_name: str | None │ +│ query: str | None (for search tools) │ +│ tool_input: str | None (legacy fallback) │ +│ reason: str │ +├──────────────────────────────────────────────────────┤ +│ CodeFindings │ +│ findings: list[str] (3-5 technical facts) │ +│ tool_outcome: str ("CALLED: X with Y → Z") │ +├──────────────────────────────────────────────────────┤ +│ Observation │ +│ results: list[str] (current step's facts) │ +│ memory: list[str] (cumulative findings) │ +├──────────────────────────────────────────────────────┤ +│ QuestionRouting │ +│ agent_type: "reachability" | "code_understanding" │ +│ reason: str │ +├──────────────────────────────────────────────────────┤ +│ Classification │ +│ is_reachability: "yes" | "no" │ +├──────────────────────────────────────────────────────┤ +│ PackageSelection │ +│ selected_package: str │ +│ reason: str │ +└──────────────────────────────────────────────────────┘ +``` + +--- + +### Memory System + +**File:** `src/vuln_analysis/functions/react_internals.py` + +The `InternalMemory` system provides deterministic, token-budgeted storage for agent observations, replacing unbounded LLM-managed memory lists. + +#### Structure + +```python +@dataclass +class MemoryItem: + key: str # Dedup key (stable across line number changes) + category: str # One of 17 categories + content: str # Full text of finding + priority: int # 1-3 (higher = keep longer) + occurrence_count: int # Dedup counter +``` + +#### Category Taxonomy + +Categories are organized by agent layer and purpose, each with a priority level: + +| Category | Priority | Context | +|---|---|---| +| `fix_found`, `fix_applied` | 3 | Patch/fix evidence | +| `fix_definition`, `fix_absent` | 2 | Patch presence/absence | +| `vuln_found` | 3 | Vulnerable code confirmed | +| `vuln_absent` | 2 | Vulnerable code not found | +| `file_compiled`, `file_not_compiled` | 3 | Build/compilation status | +| `feature_disabled`, `feature_enabled` | 3/2 | Feature flag status | +| `reachable`, `not_reachable` | 3 | Call chain results | +| `validated` | 2 | Package/function validated | +| `tool_record` | 2 | What was already tried | +| `coverage` | 1 | Coverage tracking | +| `failed` | 1 | Tool failures | +| `ambiguous` | 1 | Inconclusive findings | +| `other` | 1 | Unrecognized formats | + +#### Eviction Policy + +When memory exceeds its token budget (default 1550 tokens): +1. Find the **lowest-priority** item +2. Among equal priority, evict **oldest first** (FIFO) +3. Repeat until under budget + +This ensures high-priority findings (reachability results, vulnerability evidence) survive while low-priority items (tool failures, coverage checks) are evicted first. + +#### Deduplication + +Items are deduplicated by a key derived from their category and content. For location-sensitive categories, line numbers are stripped from the key so that the same function found at slightly different line numbers still deduplicates: + +``` +"fix_found:PQescapeLiteral@cls_api.c:1956" → key: "fix_found:PQescapeLiteral@cls_api.c" +"fix_found:PQescapeLiteral@cls_api.c:1960" → key: "fix_found:PQescapeLiteral@cls_api.c" + (deduplicates ↑) +``` + +--- + +### File Reference + +| File | Role | +|---|---| +| `functions/cve_agent.py` | Orchestrator entrypoint — fans out across CVEs and questions | +| `functions/base_graph_agent.py` | Shared state machine template (Template Method pattern) | +| `functions/reachability_agent.py` | Reachability sub-agent implementation | +| `functions/code_understanding_agent.py` | Code understanding sub-agent implementation | +| `functions/code_understanding_internals.py` | CU-specific rules tracker and prompts | +| `functions/react_internals.py` | Shared schemas, rules, prompts, memory system | +| `functions/agent_registry.py` | Agent registration and lookup | +| `functions/dispatcher.py` | LLM-based question routing | +| `tools/tool_names.py` | Tool name constants | +| `utils/prompt_factory.py` | Ecosystem-specific tool strategies (reachability) | +| `utils/code_understanding_prompt_factory.py` | Ecosystem-specific tool strategies (CU) | + +--- + +## File Reference + +| File | Role | +|---|---| +| `functions/cve_agent.py` | Orchestrator — pre-process, fan-out, post-process | +| `functions/dispatcher.py` | LLM-based question routing | +| `functions/agent_registry.py` | Sub-agent registration and lookup | +| `functions/base_graph_agent.py` | Shared five-node state machine template | +| `functions/reachability_agent.py` | Reachability / transitive search sub-agent | +| `functions/code_understanding_agent.py` | Code understanding sub-agent | +| `functions/react_internals.py` | Shared schemas, rules, prompts, memory system |