add a define context phase

src/agents/auditor/agent.ts

@@ -1,35 +1,121 @@
+import fs from "node:fs";
+import path from "node:path";
+
 import { HumanMessage, SystemMessage } from "@langchain/core/messages";
 import { END, type GraphNode, START, StateGraph } from "@langchain/langgraph";
 import { z } from "zod";
 
 import { auditorModel } from "./model.ts";
-import {
-  CRITIC_FINDINGS_PROMPT,
-  FIND_VULNERABILITIES_PROMPT,
-  GATHER_CONTEXT_PROMPT,
-} from "./prompts.ts";
+import { CRITIC_FINDINGS_PROMPT, FIND_VULNERABILITIES_PROMPT, GATHER_CONTEXT_PROMPT } from "./prompts.ts";
 import { AuditorState, CriticSchema, FindingSchema } from "./state.ts";
+import { analyzeSolidityFile } from "./tools/solidity-analyzer-tool.ts";
+import {
+  DOC_BASENAMES,
+  DOC_EXTS,
+  MAX_DEPTH,
+  MAX_DOC_CHARS,
+  MAX_REFLECTIONS,
+  MAX_SOL_CHARS,
+  SKIP_DIRS,
+  SOL_EXT,
+} from "./config.ts";
+
+const walkDirectory = (dir: string, depth: number, solFiles: string[], docFiles: string[]) => {
+  if (depth > MAX_DEPTH) return;
+
+  let entries: fs.Dirent[];
+  try {
+    entries = fs.readdirSync(dir, { withFileTypes: true });
+  } catch {
+    return;
+  }
+
+  for (const entry of entries) {
+    if (entry.isDirectory()) {
+      if (!SKIP_DIRS.has(entry.name)) {
+        walkDirectory(path.join(dir, entry.name), depth + 1, solFiles, docFiles);
+      }
+    } else if (entry.isFile()) {
+      const fullPath = path.join(dir, entry.name);
+      const ext = path.extname(entry.name).toLowerCase();
+      const base = path.basename(entry.name, ext).toLowerCase();
+
+      if (ext === SOL_EXT) {
+        solFiles.push(fullPath);
+      } else if (DOC_EXTS.has(ext) || DOC_BASENAMES.has(base)) {
+        docFiles.push(fullPath);
+      }
+    }
+  }
+};
 
-const MAX_REFLECTIONS = 3;
+const defineScope: GraphNode<typeof AuditorState> = async (state) => {
+  const solFiles: string[] = [];
+  const docFiles: string[] = [];
+
+  walkDirectory(state.repoPath, 0, solFiles, docFiles);
+
+  return { scope: solFiles, docs: docFiles };
+};
 
 const gatherContext: GraphNode<typeof AuditorState> = async (state) => {
-  const model = auditorModel.withStructuredOutput(z.object({ context: z.string() }));
+  const readFile = (filePath: string): string => {
+    try {
+      return fs.readFileSync(filePath, "utf-8");
+    } catch {
+      return "";
+    }
+  };
 
-  const scopeNote =
-    state.scope.length > 0 ? `\nAudit scope (focus on these): ${state.scope.join(", ")}` : "";
+  // Read and analyze each Solidity file
+  const solidityEntries: { filePath: string; source: string; analysis: string }[] = [];
+  for (const filePath of state.scope) {
+    const source = readFile(filePath).slice(0, MAX_SOL_CHARS);
+    if (!source) continue;
+    const analysis = await analyzeSolidityFile(source, "short");
+    solidityEntries.push({ filePath, source, analysis });
+  }
 
-  const result = await model.invoke([
-    new SystemMessage(GATHER_CONTEXT_PROMPT),
-    new HumanMessage(`Analyze this smart contract:${scopeNote}\n\n${state.solidityFile}`),
-  ]);
+  // Concatenate all sources for downstream vulnerability phases
+  const solidityFile = solidityEntries
+    .map(({ filePath, source }) => `// === FILE: ${filePath} ===\n${source}`)
+    .join("\n\n");
+
+  // Read documentation files
+  const docEntries: { filePath: string; content: string }[] = [];
+  for (const filePath of state.docs) {
+    const content = readFile(filePath).slice(0, MAX_DOC_CHARS);
+    if (content) docEntries.push({ filePath, content });
+  }
+
+  // Build the LLM input
+  const parts: string[] = [];
 
-  return { repoContext: result.context };
+  if (docEntries.length > 0) {
+    parts.push("## Documentation\n");
+    for (const { filePath, content } of docEntries) {
+      parts.push(`### ${filePath}\n${content}`);
+    }
+  }
+
+  parts.push("## Structural Analysis (auto-generated)\n");
+  for (const { filePath, analysis } of solidityEntries) {
+    parts.push(`### ${filePath}\n${analysis}`);
+  }
+
+  parts.push("## Contract Source Code\n");
+  for (const { filePath, source } of solidityEntries) {
+    parts.push(`### ${filePath}\n\`\`\`solidity\n${source}\n\`\`\``);
+  }
+
+  const model = auditorModel.withStructuredOutput(z.object({ context: z.string() }));
+  const result = await model.invoke([new SystemMessage(GATHER_CONTEXT_PROMPT), new HumanMessage(parts.join("\n\n"))]);
+
+  return { solidityFile, repoContext: result.context };
 };
 
 const findVulnerabilities: GraphNode<typeof AuditorState> = async (state) => {
-  const model = auditorModel.withStructuredOutput(
-    z.object({ findings: z.array(FindingSchema) }),
-  );
+  const model = auditorModel.withStructuredOutput(z.object({ findings: z.array(FindingSchema) }));
 
   let userMessage = `Contract:\n\n${state.solidityFile}\n\nProtocol Context:\n${state.repoContext}`;
 
@@ -43,10 +129,7 @@ const findVulnerabilities: GraphNode<typeof AuditorState> = async (state) => {
     userMessage += `\n\nCritic feedback from previous iteration (iteration ${state.reflectionCount}):\n${feedback}\n\nRevise your findings accordingly.`;
   }
 
-  const result = await model.invoke([
-    new SystemMessage(FIND_VULNERABILITIES_PROMPT),
-    new HumanMessage(userMessage),
-  ]);
+  const result = await model.invoke([new SystemMessage(FIND_VULNERABILITIES_PROMPT), new HumanMessage(userMessage)]);
 
   return { candidateFindings: result.findings };
 };
@@ -60,9 +143,7 @@ const criticFindings: GraphNode<typeof AuditorState> = async (state) => {
     };
   }
 
-  const model = auditorModel.withStructuredOutput(
-    z.object({ reviews: z.array(CriticSchema) }),
-  );
+  const model = auditorModel.withStructuredOutput(z.object({ reviews: z.array(CriticSchema) }));
 
   const findingsText = state.candidateFindings
     .map(
@@ -93,10 +174,12 @@ const criticFindings: GraphNode<typeof AuditorState> = async (state) => {
 };
 
 export const auditorAgent = new StateGraph(AuditorState)
+  .addNode("defineScope", defineScope)
   .addNode("gatherContext", gatherContext)
   .addNode("findVulnerabilities", findVulnerabilities)
   .addNode("criticFindings", criticFindings)
-  .addEdge(START, "gatherContext")
+  .addEdge(START, "defineScope")
+  .addEdge("defineScope", "gatherContext")
   .addEdge("gatherContext", "findVulnerabilities")
   .addEdge("findVulnerabilities", "criticFindings")
   .addConditionalEdges("criticFindings", (state) => {

src/agents/auditor/config.ts

@@ -0,0 +1,18 @@
+export const SOL_EXT = ".sol";
+export const DOC_EXTS = new Set([".md", ".txt", ".rst", ".adoc"]);
+export const DOC_BASENAMES = new Set(["readme", "whitepaper", "spec", "architecture", "design", "overview", "docs"]);
+export const SKIP_DIRS = new Set([
+  "node_modules",
+  ".git",
+  "out",
+  "artifacts",
+  "cache",
+  "lib",
+  ".deps",
+  "build",
+  "dist",
+]);
+export const MAX_DEPTH = 6;
+export const MAX_DOC_CHARS = 12_000;
+export const MAX_SOL_CHARS = 40_000;
+export const MAX_REFLECTIONS = 3;

src/agents/auditor/prompts.ts

@@ -1,14 +1,24 @@
-export const GATHER_CONTEXT_PROMPT = `You are a smart contract security expert. Analyze the provided Solidity source code and produce a thorough protocol context that will guide vulnerability discovery.
+export const GATHER_CONTEXT_PROMPT = `You are a smart contract security expert. You will receive documentation, a structural analysis, and the full source of all in-scope Solidity contracts. Produce a thorough protocol context that will guide vulnerability discovery.
 
-Extract and clearly structure the following:
+Structure your output in the following sections:
 
-1. **Invariants**: Conditions that must always hold (e.g., "total supply must equal sum of all balances", "contract ETH balance >= sum of all user deposits").
+## 1. Contract Overview
+For each contract: its purpose, kind (contract/interface/library/abstract), inheritance chain, and key dependencies on other in-scope contracts or external protocols.
 
-2. **Design Assumptions**: What the protocol assumes about callers, external contracts, oracles, admin keys, and token behavior (e.g., "tokens are ERC-20 compliant", "admin is trusted", "no fee-on-transfer tokens").
+## 2. State & Storage Map
+List all meaningful state variables across contracts, what they represent, and which functions read or write them. Flag shared or inherited storage.
 
-3. **Key Flows**: The main execution paths and state transitions (e.g., deposit → mint shares → updateRewards; withdraw → burn shares → transfer ETH).
+## 3. Key Flows
+Trace the main execution paths and state transitions end-to-end across contracts (e.g., deposit → mint shares → updateRewards; withdraw → burn shares → transfer ETH). Include cross-contract calls.
 
-4. **Business Rules**: Access controls, fee structures, timelocks, caps, pausing mechanisms, and any other domain constraints.
+## 4. Invariants
+Conditions that must always hold (e.g., "total supply must equal sum of all balances", "contract ETH balance ≥ sum of all user deposits"). Derive these from both the source and any documentation.
+
+## 5. Design Assumptions
+What the protocol assumes about callers, external contracts, oracles, admin keys, and token behavior (e.g., "tokens are ERC-20 compliant", "admin is trusted", "no fee-on-transfer tokens").
+
+## 6. Business Rules
+Access controls, fee structures, timelocks, caps, pausing mechanisms, upgrade patterns, and any other domain constraints.
 
 Be precise and exhaustive — the richer the context, the more accurately vulnerabilities can be identified and validated.`;
 

src/agents/auditor/state.ts

@@ -22,8 +22,10 @@ export const CriticSchema = z.object({
 });
 
 export const AuditorState = new StateSchema({
+  repoPath: z.string().default(""),
   solidityFile: z.string().default(""),
   scope: z.array(z.string()).default([]),
+  docs: z.array(z.string()).default([]),
   repoContext: z.string().default(""),
   candidateFindings: z.array(FindingSchema).default([]),
   criticReviews: z.array(CriticSchema).default([]),

src/agents/auditor/tools/solidity-analyzer-tool.ts

@@ -434,7 +434,7 @@ const generateMarkdown = (imports: string[], contracts: ContractAnalysis[]) => {
   return lines.join("\n");
 };
 
-const analyzeSolidityFile = async (soliditySource: string, mode: "full" | "short") => {
+export const analyzeSolidityFile = async (soliditySource: string, mode: "full" | "short") => {
   let ast: any;
 
   try {