adjust audit schemas

src/agents/auditor/agent.ts

@@ -8,7 +8,7 @@ import { z } from "zod";
 import { logger } from "../../logger.ts";
 import { judgeFindingsModel, findVulnerabilitiesModel, gatherContextModel } from "./model.ts";
 import { JUDGE_FINDINGS_PROMPT, FIND_VULNERABILITIES_PROMPT, GATHER_CONTEXT_PROMPT } from "./prompts.ts";
-import { AuditorState, ReviewSchema, PartialFindingSchema } from "./state.ts";
+import { AuditorState, JudgeReviewSchema, CandidateFindingSchema } from "./state.ts";
 import { analyzeSolidityFile } from "./tools/solidity-analyzer-tool.ts";
 import { buildRepoTree } from "./tools/repo-tree-tool.ts";
 import {
@@ -91,11 +91,6 @@ const gatherContext: GraphNode<typeof AuditorState> = async (state) => {
     solidityEntries.push({ filePath, source, analysis });
   }
 
-  // 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) {
@@ -123,38 +118,61 @@ const gatherContext: GraphNode<typeof AuditorState> = async (state) => {
     parts.push(`### ${filePath}\n\`\`\`solidity\n${source}\n\`\`\``);
   }
 
-  // const model = gatherContextModel.withStructuredOutput(z.object({ context: z.string() }));
-  // const result = await model.invoke([new SystemMessage(GATHER_CONTEXT_PROMPT), new HumanMessage(parts.join("\n\n"))]);
+  const model = gatherContextModel.withStructuredOutput(z.object({ context: z.string() }));
+  const result = await model.invoke([new SystemMessage(GATHER_CONTEXT_PROMPT), new HumanMessage(parts.join("\n\n"))]);
 
   logger.info(`gatherContext: context built (${parts.join("\n\n").length} chars)`);
   logger.debug(`gatherContext: full context:\n${parts.join("\n\n")}`);
 
-  return { solidityFile, repoContext: parts.join("\n\n") };
+  return { repoContext: result.context };
 };
 
 const findVulnerabilities: GraphNode<typeof AuditorState> = async (state) => {
-  const model = findVulnerabilitiesModel.withStructuredOutput(z.object({ findings: z.array(PartialFindingSchema) }));
-
-  let userMessage = `Contract:\n\n${state.solidityFile}\n\nProtocol Context:\n${state.repoContext}`;
-
-  if (state.judgeReviews.length > 0) {
-    const feedback = state.judgeReviews
-      .map(
-        (r) => `- "${r /*.title*/}": ${r.isFalsePositive ? "FALSE POSITIVE" : "TRUE POSITIVE"}\n  Judge: ${r.review}`,
-      )
-      .join("\n");
-    userMessage += `\n\nJudge feedback from previous iteration (iteration ${state.reflectionCount}):\n${feedback}\n\nRevise your findings accordingly.`;
-  }
+  const model = findVulnerabilitiesModel.withStructuredOutput(z.object({ findings: z.array(CandidateFindingSchema) }));
+
+  const previousFeedback =
+    state.judgeReviews.length > 0
+      ? state.judgeReviews
+          .map((r, i) => {
+            const title = state.candidateFindings[i]?.title ?? `Finding ${i + 1}`;
+            return `- "${title}": ${r.isFalsePositive ? "FALSE POSITIVE" : "TRUE POSITIVE"}\n  Judge: ${r.review}`;
+          })
+          .join("\n")
+      : null;
+
+  logger.info(
+    `findVulnerabilities: invoking LLM for ${state.scope.length} file(s) in parallel (iteration ${state.reflectionCount + 1})`,
+  );
+
+  const allFindings = await Promise.all(
+    state.scope.map(async (filePath) => {
+      let source: string;
+      try {
+        source = fs.readFileSync(filePath, "utf-8").slice(0, MAX_SOL_CHARS);
+      } catch {
+        return [];
+      }
+      if (!source) return [];
 
-  logger.info(`findVulnerabilities: invoking LLM (iteration ${state.reflectionCount + 1})`);
-  logger.debug(`findVulnerabilities: user message:\n${userMessage}`);
+      let userMessage = `Contract (${filePath}):\n\n${source}\n\nProtocol Context:\n${state.repoContext}`;
+      if (previousFeedback) {
+        userMessage += `\n\nJudge feedback from previous iteration (iteration ${state.reflectionCount}):\n${previousFeedback}\n\nRevise your findings accordingly.`;
+      }
 
-  const result = await model.invoke([new SystemMessage(FIND_VULNERABILITIES_PROMPT), new HumanMessage(userMessage)]);
+      logger.debug(`findVulnerabilities: processing ${filePath}`);
+      const result = await model.invoke([
+        new SystemMessage(FIND_VULNERABILITIES_PROMPT),
+        new HumanMessage(userMessage),
+      ]);
+      return result.findings.map((finding: any) => ({ ...finding, path: filePath, location: "1-14" }));
+    }),
+  );
 
-  logger.info(`findVulnerabilities: LLM returned ${result.findings.length} candidate finding(s)`);
-  logger.debug(`findVulnerabilities: findings:\n${JSON.stringify(result.findings, null, 2)}`);
+  const candidateFindings = allFindings.flat();
+  logger.info(`findVulnerabilities: LLM returned ${candidateFindings.length} total candidate finding(s)`);
+  logger.debug(`findVulnerabilities: findings:\n${JSON.stringify(candidateFindings, null, 2)}`);
 
-  return { candidateFindings: result.findings.map((finding: any) => ({ ...finding, path: "/", location: "1-14" })) };
+  return { candidateFindings };
 };
 
 const judgeFindings: GraphNode<typeof AuditorState> = async (state) => {
@@ -167,39 +185,52 @@ const judgeFindings: GraphNode<typeof AuditorState> = async (state) => {
     };
   }
 
-  const model = findVulnerabilitiesModel.withStructuredOutput(z.object({ reviews: z.array(ReviewSchema) }));
+  const model = judgeFindingsModel.withStructuredOutput(JudgeReviewSchema);
+
+  logger.info(`judgeFindings: reviewing ${state.candidateFindings.length} candidate finding(s) in parallel`);
+
+  const reviews = await Promise.all(
+    state.candidateFindings.map(async (finding, i) => {
+      let source: string;
+      try {
+        source = fs.readFileSync(finding.path, "utf-8").slice(0, MAX_SOL_CHARS);
+      } catch {
+        source = "";
+      }
 
-  const findingsText = state.candidateFindings
-    .map(
-      (f, i) =>
-        `[Finding ${i + 1}] ${f.title}\nSeverity: ${f.severity}\nDescription: ${f.description}\nLocation: ${f.path} lines ${f.location}\nCode:\n\`\`\`solidity\n${f.codeSnippet}\n\`\`\``,
-    )
-    .join("\n\n---\n\n");
+      const findingText = `[Finding ${i + 1}] ${finding.title}\nSeverity: ${finding.severity}\nDescription: ${finding.description}\nLocation: ${finding.path} lines ${finding.location}\nCode:\n\`\`\`solidity\n${finding.codeSnippet}\n\`\`\``;
 
-  logger.info(`judgeFindings: reviewing ${state.candidateFindings.length} candidate finding(s)`);
-  logger.debug(`judgeFindings: findings text:\n${findingsText}`);
+      logger.debug(`judgeFindings: reviewing finding ${i + 1}: ${finding.title}`);
+      return model.invoke([
+        new SystemMessage(JUDGE_FINDINGS_PROMPT),
+        new HumanMessage(
+          `Contract (${finding.path}):\n\n${source}\n\nProtocol Context:\n${state.repoContext}\n\nFinding to Review:\n\n${findingText}`,
+        ),
+      ]);
+    }),
+  );
 
-  const result = await model.invoke([
-    new SystemMessage(JUDGE_FINDINGS_PROMPT),
-    new HumanMessage(
-      `Contract:\n\n${state.solidityFile}\n\nProtocol Context:\n${state.repoContext}\n\nCandidate Findings to Review:\n\n${findingsText}`,
-    ),
-  ]);
+  const confirmedEntries = state.candidateFindings
+    .map((finding, i) => ({ finding, review: reviews[i] }))
+    .filter(({ review }) => !review.isFalsePositive);
 
-  const reviewsByTitle = new Map(result.reviews.map((r: any) => [r.findingTitle.toLowerCase(), r]));
+  const findings = confirmedEntries.map(({ finding, review }) => ({
+    ...finding,
+    judgeReview: {
+      review: review.review,
+      confidence: review.confidence,
+      exploitablePaths: review.exploitablePaths,
+    },
+  }));
 
-  const confirmedFindings = state.candidateFindings.filter((f, i) => {
-    const review = reviewsByTitle.get(f.title.toLowerCase()) ?? result.reviews[i];
-    return review ? !review.isFalsePositive : true;
-  });
+  const falsePositiveCount = state.candidateFindings.length - findings.length;
 
-  const falsePositiveCount = state.candidateFindings.length - confirmedFindings.length;
-  logger.info(`judgeFindings: ${confirmedFindings.length} confirmed, ${falsePositiveCount} false positive(s)`);
-  logger.debug(`judgeFindings: reviews:\n${JSON.stringify(result.reviews, null, 2)}`);
+  logger.info(`judgeFindings: ${findings.length} confirmed, ${falsePositiveCount} false positive(s)`);
+  logger.debug(`judgeFindings: reviews:\n${JSON.stringify(reviews, null, 2)}`);
 
   return {
-    judgeReviews: result.reviews,
-    findings: confirmedFindings,
+    judgeReviews: reviews,
+    findings,
     reflectionCount: state.reflectionCount + 1,
   };
 };

src/agents/auditor/config.ts

@@ -20,4 +20,4 @@ export const SKIP_DIRS = new Set([
 export const MAX_DEPTH = 6;
 export const MAX_DOC_CHARS = 12_000;
 export const MAX_SOL_CHARS = 40_000;
-export const MAX_REFLECTIONS = 3;
+export const MAX_REFLECTIONS = 1;

src/agents/auditor/model.ts

@@ -5,9 +5,9 @@ export const gatherContextModel = new ChatAnthropic({
 });
 
 export const findVulnerabilitiesModel = new ChatAnthropic({
-  model: "claude-sonnet-4-6",
+  model: "claude-haiku-4-5",
 });
 
 export const judgeFindingsModel = new ChatAnthropic({
-  model: "claude-sonnet-4-6",
+  model: "claude-haiku-4-5",
 });

src/agents/auditor/state.ts

@@ -1,7 +1,7 @@
 import { StateSchema } from "@langchain/langgraph";
 import { z } from "zod";
 
-export const PartialFindingSchema = z.object({
+export const CandidateFindingSchema = z.object({
   title: z.string(),
   description: z.string(),
   recommendation: z.string(),
@@ -9,32 +9,34 @@ export const PartialFindingSchema = z.object({
   codeSnippet: z.string(),
 });
 
-export const FindingSchema = z.object({
-  title: z.string(),
-  description: z.string(),
-  recommendation: z.string(),
-  severity: z.enum(["high", "medium", "low"]),
-  codeSnippet: z.string(),
+export const LocatedFindingSchema = CandidateFindingSchema.extend({
   location: z.string(),
   path: z.string(),
 });
 
-export const ReviewSchema = z.object({
+export const JudgeReviewSchema = z.object({
   review: z.string(),
   isFalsePositive: z.boolean(),
   confidence: z.number(),
   exploitablePaths: z.array(z.string()),
 });
 
+export const FindingSchema = LocatedFindingSchema.extend({
+  judgeReview: z.object({
+    review: z.string(),
+    confidence: z.number(),
+    exploitablePaths: z.array(z.string()),
+  }),
+});
+
 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([]),
   fileTree: z.string().default(""),
   repoContext: z.string().default(""),
-  candidateFindings: z.array(FindingSchema).default([]),
-  judgeReviews: z.array(ReviewSchema).default([]),
+  candidateFindings: z.array(LocatedFindingSchema).default([]),
+  judgeReviews: z.array(JudgeReviewSchema).default([]),
   findings: z.array(FindingSchema).default([]),
   reflectionCount: z.number().default(0),
 });

src/index.ts

@@ -1,22 +1,35 @@
-import "dotenv/config";
+// import "dotenv/config";
 
-import { auditorAgent } from "./agents/auditor/agent.ts";
-import { coderAgent } from "./agents/coder/agent.ts";
-import { testerAgent } from "./agents/tester/agent.ts";
+// import { auditorAgent } from "./agents/auditor/agent.ts";
+// import { coderAgent } from "./agents/coder/agent.ts";
+// import { testerAgent } from "./agents/tester/agent.ts";
+
+// const requirements = ["ERC20 token", "pausable", "ownable"];
+
+// const coderResult = await coderAgent.invoke({ requirements });
+// console.log("======= Coder =======");
+// console.log(coderResult.contract);
 
-const requirements = ["ERC20 token", "pausable", "ownable"];
+// const auditorResult = await auditorAgent.invoke({ solidityFile: coderResult.contract });
+// console.log("\n======= Auditor =======");
+// console.log(auditorResult.findings);
 
-const coderResult = await coderAgent.invoke({ requirements });
-console.log("======= Coder =======");
-console.log(coderResult.contract);
+// const testerResult = await testerAgent.invoke({
+//   solidityFiles: [coderResult.contract],
+//   vulnerability: auditorResult.findings[0] ?? {},
+// });
+// console.log("\n======= Tester =======");
+// console.log(testerResult.results);
 
-const auditorResult = await auditorAgent.invoke({ solidityFile: coderResult.contract });
-console.log("\n======= Auditor =======");
-console.log(auditorResult.findings);
+import "dotenv/config";
+import { auditorAgent } from "./agents/auditor/agent.ts";
+import { logger } from "./logger.ts";
 
-const testerResult = await testerAgent.invoke({
-  solidityFiles: [coderResult.contract],
-  vulnerability: auditorResult.findings[0] ?? {},
+logger.info("Starting auditorAgent");
+
+const result = await auditorAgent.invoke({
+  repoPath: "/Users/uanderson/personal/projeto-talp1/dist/repos/example",
 });
-console.log("\n======= Tester =======");
-console.log(testerResult.results);
+
+logger.info("Agent completed");
+logger.debug(`Agent result:\n${JSON.stringify(result, null, 2)}`);