add agent phases

src/agents/auditor/agent.ts

@@ -1,20 +1,109 @@
+import { HumanMessage, SystemMessage } from "@langchain/core/messages";
 import { END, type GraphNode, START, StateGraph } from "@langchain/langgraph";
+import { z } from "zod";
 
-import { AuditorState } from "./state.ts";
-import { slitherTool } from "./tools/slither-tool.ts";
+import { auditorModel } from "./model.ts";
+import {
+  CRITIC_FINDINGS_PROMPT,
+  FIND_VULNERABILITIES_PROMPT,
+  GATHER_CONTEXT_PROMPT,
+} from "./prompts.ts";
+import { AuditorState, CriticSchema, FindingSchema } from "./state.ts";
 
-const PLACEHOLDER_VULNERABILITIES = [
-  { type: "reentrancy", severity: "high", description: "Unchecked external call allows reentrancy attack." },
-  { type: "integer-overflow", severity: "medium", description: "Arithmetic operation may overflow." },
-];
+const MAX_REFLECTIONS = 3;
 
-const auditContract: GraphNode<typeof AuditorState> = async (state) => {
-  await slitherTool.invoke({ solidityFile: state.solidityFile });
-  return { vulnerabilities: PLACEHOLDER_VULNERABILITIES };
+const gatherContext: GraphNode<typeof AuditorState> = async (state) => {
+  const model = auditorModel.withStructuredOutput(z.object({ context: z.string() }));
+
+  const scopeNote =
+    state.scope.length > 0 ? `\nAudit scope (focus on these): ${state.scope.join(", ")}` : "";
+
+  const result = await model.invoke([
+    new SystemMessage(GATHER_CONTEXT_PROMPT),
+    new HumanMessage(`Analyze this smart contract:${scopeNote}\n\n${state.solidityFile}`),
+  ]);
+
+  return { repoContext: result.context };
+};
+
+const findVulnerabilities: GraphNode<typeof AuditorState> = async (state) => {
+  const model = auditorModel.withStructuredOutput(
+    z.object({ findings: z.array(FindingSchema) }),
+  );
+
+  let userMessage = `Contract:\n\n${state.solidityFile}\n\nProtocol Context:\n${state.repoContext}`;
+
+  if (state.criticReviews.length > 0) {
+    const feedback = state.criticReviews
+      .map(
+        (r) =>
+          `- "${r.findingTitle}": ${r.isFalsePositive ? "FALSE POSITIVE" : "TRUE POSITIVE"}\n  Critic: ${r.review}`,
+      )
+      .join("\n");
+    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),
+  ]);
+
+  return { candidateFindings: result.findings };
+};
+
+const criticFindings: GraphNode<typeof AuditorState> = async (state) => {
+  if (state.candidateFindings.length === 0) {
+    return {
+      criticReviews: [],
+      findings: [],
+      reflectionCount: state.reflectionCount + 1,
+    };
+  }
+
+  const model = auditorModel.withStructuredOutput(
+    z.object({ reviews: z.array(CriticSchema) }),
+  );
+
+  const findingsText = state.candidateFindings
+    .map(
+      (f, i) =>
+        `[Finding ${i + 1}] ${f.title}\nSeverity: ${f.severity} | Auditor confidence: ${f.confidence}/100\nDescription: ${f.description}\nLocation: ${f.path} lines ${f.location}\nCode:\n\`\`\`solidity\n${f.codeSnippet}\n\`\`\`\nExploit paths:\n${f.exploitablePaths.map((p) => `  - ${p}`).join("\n")}`,
+    )
+    .join("\n\n---\n\n");
+
+  const result = await model.invoke([
+    new SystemMessage(CRITIC_FINDINGS_PROMPT),
+    new HumanMessage(
+      `Contract:\n\n${state.solidityFile}\n\nProtocol Context:\n${state.repoContext}\n\nCandidate Findings to Review:\n\n${findingsText}`,
+    ),
+  ]);
+
+  const reviewsByTitle = new Map(result.reviews.map((r) => [r.findingTitle.toLowerCase(), r]));
+
+  const confirmedFindings = state.candidateFindings.filter((f, i) => {
+    const review = reviewsByTitle.get(f.title.toLowerCase()) ?? result.reviews[i];
+    return review ? !review.isFalsePositive : true;
+  });
+
+  return {
+    criticReviews: result.reviews,
+    findings: confirmedFindings,
+    reflectionCount: state.reflectionCount + 1,
+  };
 };
 
 export const auditorAgent = new StateGraph(AuditorState)
-  .addNode("auditContract", auditContract)
-  .addEdge(START, "auditContract")
-  .addEdge("auditContract", END)
+  .addNode("gatherContext", gatherContext)
+  .addNode("findVulnerabilities", findVulnerabilities)
+  .addNode("criticFindings", criticFindings)
+  .addEdge(START, "gatherContext")
+  .addEdge("gatherContext", "findVulnerabilities")
+  .addEdge("findVulnerabilities", "criticFindings")
+  .addConditionalEdges("criticFindings", (state) => {
+    const hasFalsePositives = state.criticReviews.some((r) => r.isFalsePositive);
+    if (hasFalsePositives && state.reflectionCount < MAX_REFLECTIONS) {
+      return "findVulnerabilities";
+    }
+    return END;
+  })
   .compile();

src/agents/auditor/prompts.ts

@@ -0,0 +1,45 @@
+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.
+
+Extract and clearly structure the following:
+
+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").
+
+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").
+
+3. **Key Flows**: The main execution paths and state transitions (e.g., deposit → mint shares → updateRewards; withdraw → burn shares → transfer ETH).
+
+4. **Business Rules**: Access controls, fee structures, timelocks, caps, pausing mechanisms, and any other domain constraints.
+
+Be precise and exhaustive — the richer the context, the more accurately vulnerabilities can be identified and validated.`;
+
+export const FIND_VULNERABILITIES_PROMPT = `You are an expert smart contract security auditor specializing in Solidity. Systematically analyze the contract source code and protocol context to identify security vulnerabilities.
+
+For each vulnerability provide ALL of the following fields:
+
+- **title**: Short, precise name (e.g., "Reentrancy in withdraw", "Missing access control on setFee").
+- **description**: Explain the EXPECTED behavior vs the OBSERVED (vulnerable) behavior in 2–4 sentences.
+- **recommendation**: Specific, actionable remediation (e.g., "Apply checks-effects-interactions pattern", "Add onlyOwner modifier").
+- **severity**: One of "high" (direct fund loss or contract takeover), "medium" (indirect or conditional risk), "low" (best-practice issue, no immediate financial risk).
+- **confidence**: Integer 0–100 reflecting your confidence this is a real, exploitable vulnerability.
+- **codeSnippet**: The exact vulnerable code block as it appears in the source.
+- **location**: Line range as "start-end" (e.g., "42-58"). Use "unknown" if lines are not determinable.
+- **path**: File path of the vulnerable contract. Use the contract name if a single inline source is provided.
+- **exploitablePaths**: Array of one or more concrete exploit traces. Each trace must describe the attacker steps with realistic inputs/values (e.g., "1. Attacker calls deposit(100 ETH) 2. Attacker contract fallback re-enters withdraw() before balance update 3. Attacker drains 100 ETH twice").
+
+Vulnerability categories to systematically check: reentrancy (single- and cross-function), access control, integer overflow/underflow, oracle manipulation, flash loan attacks, front-running/MEV, signature replay, storage collisions, uninitialized proxies, unsafe delegatecall, gas griefing, denial of service, precision loss, and logic/business rule violations.
+
+If critic feedback is provided from a previous iteration, remove confirmed false positives from your list and refine or expand remaining findings based on the critique.`;
+
+export const CRITIC_FINDINGS_PROMPT = `You are a rigorous smart contract security reviewer. Evaluate each candidate vulnerability submitted by the auditor and determine whether it is a true positive or a false positive.
+
+For each finding provide ALL of the following fields:
+
+- **findingTitle**: Must match exactly the title of the finding you are reviewing.
+- **review**: Detailed analysis (3–6 sentences) explaining why the vulnerability is or isn't real. Reference specific code, protocol invariants, preconditions, and mitigating controls.
+- **isFalsePositive**: true if the finding is NOT exploitable in practice; false if it IS a real vulnerability.
+- **confidence**: Integer 0–100 reflecting your confidence in this verdict.
+- **exploitablePaths**: If a true positive, provide concrete paths confirming exploitability with real values. If a false positive, provide the reasoning that blocks the exploit.
+
+A finding is a false positive if and only if: the exploit path is unreachable given access controls or preconditions, it is already fully mitigated by the code, it requires impossible or economically infeasible conditions, or it is explicitly documented as by-design behavior in the protocol assumptions.
+
+You must provide exactly one review object per finding, in the same order as the findings were presented.`;

src/agents/auditor/state.ts

@@ -1,7 +1,32 @@
 import { StateSchema } from "@langchain/langgraph";
 import { z } from "zod";
 
+export const FindingSchema = z.object({
+  title: z.string(),
+  description: z.string(),
+  recommendation: z.string(),
+  severity: z.enum(["high", "medium", "low"]),
+  confidence: z.number(),
+  codeSnippet: z.string(),
+  location: z.string(),
+  path: z.string(),
+  exploitablePaths: z.array(z.string()),
+});
+
+export const CriticSchema = z.object({
+  findingTitle: z.string(),
+  review: z.string(),
+  isFalsePositive: z.boolean(),
+  confidence: z.number(),
+  exploitablePaths: z.array(z.string()),
+});
+
 export const AuditorState = new StateSchema({
   solidityFile: z.string().default(""),
-  vulnerabilities: z.array(z.record(z.string(), z.any())).default([]),
+  scope: z.array(z.string()).default([]),
+  repoContext: z.string().default(""),
+  candidateFindings: z.array(FindingSchema).default([]),
+  criticReviews: z.array(CriticSchema).default([]),
+  findings: z.array(FindingSchema).default([]),
+  reflectionCount: z.number().default(0),
 });

src/index.ts

@@ -12,11 +12,11 @@ console.log(coderResult.contract);
 
 const auditorResult = await auditorAgent.invoke({ solidityFile: coderResult.contract });
 console.log("\n======= Auditor =======");
-console.log(auditorResult.vulnerabilities);
+console.log(auditorResult.findings);
 
 const testerResult = await testerAgent.invoke({
   solidityFiles: [coderResult.contract],
-  vulnerability: auditorResult.vulnerabilities[0] ?? {},
+  vulnerability: auditorResult.findings[0] ?? {},
 });
 console.log("\n======= Tester =======");
 console.log(testerResult.results);