From 1b1da09ff32e2c73158894aa7dd407b21bd2ee40 Mon Sep 17 00:00:00 2001 From: Rodrigo Verdiani Date: Sat, 4 Jul 2026 18:40:21 -0300 Subject: [PATCH] feat(ai): suggestPseudocode tool, simplified prompt, DecompiledResult.found, and address sort --- .../ai/service/ChatContextBuilder.java | 71 ++++++------------- .../service/tools/FunctionMutationTools.java | 43 +++++++++++ .../ai/service/tools/FunctionTools.java | 13 ++-- 3 files changed, 73 insertions(+), 54 deletions(-) diff --git a/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java b/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java index b129bc9..3cf2682 100644 --- a/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java +++ b/src/main/java/ai/decompile/ai/service/ChatContextBuilder.java @@ -60,60 +60,31 @@ public class ChatContextBuilder { getFunctionCallGraph (or getFunctionAssembly + getCallers + getCallees) to examine the function AND the functions it calls. A thorough analysis includes understanding what callees do and how they relate to the caller. - Never just read one function in isolation and suggest analyzing callees - "later" — analyze them now as part of your response. + Never just read one function in isolation — analyze callees now. - PRIORITY RULE: When the user asks for a specific mutation (rename, comment, - signature, flag), use that tool first. Do NOT respond with a different mutation - type instead (e.g., don't suggest renames when the user asked for comments). - This rule is about mutation tool selection ONLY — it does NOT limit your - analysis. Always do a thorough analysis including callees and callers. + [MUTATION TOOLS — MANDATORY] + When your analysis reveals findings deserving a rename, comment, signature, + flag, or label change, you MUST call the corresponding tool IMMEDIATELY. + A suggestion described only in text will NOT be reflected in the UI and is + USELESS — the user needs the tool call to see and approve it. - [FUNCTION RENAME] - Use renameFunction when the user explicitly asks to rename a specific function - (applies immediately). Use suggestRenameFunction when you independently identify - a function's purpose and want to suggest a better name (requires user approval). - Proactively suggest renames whenever you identify a function's purpose with high - confidence. For example: - - A function that decrypts data with RC4 → CALL suggestRenameFunction("DecryptRC4") - - A function that parses HTTP headers → CALL suggestRenameFunction("ParseHttpHeaders") - - A function that validates a checksum → CALL suggestRenameFunction("ValidateChecksum") - - A function that is the WinMain entry point → CALL suggestRenameFunction("WinMain") - CRITICAL: When you decide a function should be renamed, you MUST call the - suggestRenameFunction tool with the function ID and new name. NEVER just - describe the rename in your text response without calling the tool. - Always call the tool, then briefly explain your reasoning afterward. + For proactive suggestions (requires user approval), use the suggest* tools: + suggestRenameFunction, suggestAddComment, suggestUpdateSignature, + suggestPseudocode, suggestRenameLabel, suggestRenameGlobalVar, + suggestRenameStruct - [FUNCTION SIGNATURE] - Use updateSignature (apply) or suggestUpdateSignature (pending approval) - to set a function's type signature. Infer parameter types and return type - from assembly analysis. Example: 'int DecryptRC4(int keyLen, const uint8_t* key)'. + For changes the user explicitly asked for (applied immediately), use: + renameFunction, addComment, updateSignature, + addFlag, removeFlag, renameLabel, + renameGlobalVar, renameStruct - [FUNCTION COMMENTS] - Use addComment (apply) or suggestAddComment (pending approval) to document - analysis findings. You can attach comments to specific addresses - (e.g., address="0x5EB5") to document individual instructions inline. - Omit address to attach to the function. The tool appends automatically. - CRITICAL: Pass ONLY the new comment text — never include previous comments. - CRITICAL: When the user asks to add comments, call addComment/suggestAddComment. - Do NOT respond with renames or signature changes unless the user also asked for those. - - [FUNCTION FLAGS] - Use addFlag and removeFlag to categorize functions. - Valid flags: CRYPTO, NETWORK, FILE_IO, ANTI_DEBUG, ENTRY_POINT, - CALLBACK, MEMORY, PROCESS, REGISTRY, UI. - - [GLOBAL VARIABLES] - Use searchGlobalVars to find data label IDs, then renameGlobalVar (apply) - or suggestRenameGlobalVar (pending approval). - - [LABELS] - Use searchLabels to find label IDs within a function, then renameLabel (apply) - or suggestRenameLabel (pending approval). - - [STRUCTS] - Use searchStructs to find struct IDs, then renameStruct (apply) - or suggestRenameStruct (pending approval). + Rules: + 1. Call mutation tools BEFORE writing your analysis text. + 2. If you have 3 suggestions, make 3 tool calls — not 1 text block. + 3. For pseudocode: only suggestPseudocode when getFunctionDecompiled returned + found=false and you are confident in the assembly logic. + 4. When the user explicitly asks for a specific change type (e.g. "add a comment"), + use that tool first. Do not respond with a different mutation type instead. %s [BINARY OVERVIEW] Filename: %s diff --git a/src/main/java/ai/decompile/ai/service/tools/FunctionMutationTools.java b/src/main/java/ai/decompile/ai/service/tools/FunctionMutationTools.java index 7ef8b7e..48077f4 100644 --- a/src/main/java/ai/decompile/ai/service/tools/FunctionMutationTools.java +++ b/src/main/java/ai/decompile/ai/service/tools/FunctionMutationTools.java @@ -2,6 +2,7 @@ package ai.decompile.ai.service.tools; import ai.decompile.analysis.model.repository.StaticFunctionRepository; import ai.decompile.analysis.service.MutationService; +import java.util.List; import java.util.UUID; import lombok.RequiredArgsConstructor; import org.springframework.ai.tool.annotation.Tool; @@ -217,6 +218,48 @@ public class FunctionMutationTools { mutationService.suggestComment(funcId, comment, address, "AI", sessionId)); } + @Tool( + description = + """ + Suggest C-like pseudocode for a function that currently has no decompiled code. + CRITICAL: Only call this when getFunctionDecompiled returns found=false and you + have thoroughly analyzed the function's assembly (call getFunctionAssembly first). + Be conservative — only suggest pseudocode when the assembly logic is clear and + you can produce accurate, idiomatic C code. Do NOT guess or fabricate logic. + The pseudocode must faithfully reproduce the exact logic shown in the assembly, + using proper C syntax (control flow, variable names derived from stack offsets, + function calls using resolved names when available). + Provide one line of code per array element. Do NOT combine multiple statements + into a single string element. A comment "Pseudocode generated by AI" will be + automatically added. + """) + public String suggestPseudocode( + @ToolParam(description = "The function ID (UUID from searchFunctions result)") + String functionId, + @ToolParam( + description = + "Array of C-like pseudocode lines. Each element is exactly ONE line of code. " + + "Do NOT put multiple statements in one element. " + + "Do not include markdown code fences.") + List pseudocodeLines) { + + if (pseudocodeLines == null || pseudocodeLines.isEmpty()) { + return error(functionId, "Pseudocode must not be empty"); + } + + var pseudocode = String.join("\n", pseudocodeLines); + var funcId = UUID.fromString(functionId); + if (functionRepository.findById(funcId).isEmpty()) { + return error(functionId, "Function not found: " + functionId); + } + + mutationService.applyComment(funcId, "Pseudocode generated by AI", "", "AI"); + + var sessionId = toolContext.getSessionId(); + return helper.formatResult( + mutationService.suggest("FUNCTION", funcId, "decompiledCode", pseudocode, "AI", sessionId)); + } + @Tool( description = """ diff --git a/src/main/java/ai/decompile/ai/service/tools/FunctionTools.java b/src/main/java/ai/decompile/ai/service/tools/FunctionTools.java index 83ed276..f16e9cd 100644 --- a/src/main/java/ai/decompile/ai/service/tools/FunctionTools.java +++ b/src/main/java/ai/decompile/ai/service/tools/FunctionTools.java @@ -19,6 +19,7 @@ import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.tool.annotation.ToolParam; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.data.domain.PageRequest; +import org.springframework.data.domain.Sort; import org.springframework.stereotype.Component; import org.springframework.transaction.annotation.Transactional; @@ -155,16 +156,18 @@ public class FunctionTools { var func = functionRepository.findById(UUID.fromString(functionId)).orElse(null); if (Objects.isNull(func)) { - return new DecompiledResult("unknown", "unknown", ""); + return new DecompiledResult("unknown", "unknown", "", false); } var renameMap = mutationService.buildRenameMap(func.getAnalysis().getId()); var code = func.getDecompiledCode(); + var hasDecompiled = Objects.nonNull(code) && !code.isBlank(); return new DecompiledResult( func.getName(), func.getAddress(), - AnalysisService.resolveNames(Objects.nonNull(code) ? code : "", renameMap)); + hasDecompiled ? AnalysisService.resolveNames(code, renameMap) : "", + hasDecompiled); } @Tool( @@ -291,7 +294,8 @@ public class FunctionTools { var page = functionRepository.findByAnalysisId( - analysis.getId(), PageRequest.of(0, MAX_LIST_FUNCTIONS)); + analysis.getId(), + PageRequest.of(0, MAX_LIST_FUNCTIONS, Sort.by("address").ascending())); return page.getContent().stream().map(this::toSummary).toList(); } @@ -345,7 +349,8 @@ public class FunctionTools { public record FunctionSummary( UUID id, String name, String address, int functionSize, String signature) {} - public record DecompiledResult(String name, String address, String decompiledCode) {} + public record DecompiledResult( + String name, String address, String decompiledCode, boolean found) {} public record FunctionAssembly(String name, String address, String assembly) {} }