开发 IDEA 版 AI 代码助手：Spring AI 与项目上下文实践

解析项目上下文

这才是插件的灵魂。我们通过 Maven API 读 pom.xml，拿到 groupId、artifactId、依赖列表；再遍历 src/main/java 目录，提取出所有包名。JavaParser 用来验证生成的代码语法是否正确。

@Service
public class ProjectContextParser {
    
    public ProjectContext parseMavenProject(String pomPath) throws Exception {
        ProjectContext context = new ProjectContext();
        File pomFile = new File(pomPath);
        MavenXpp3Reader reader = new MavenXpp3Reader();
        Model model = reader.read(new FileReader(pomFile));
        
        context.setGroupId(model.getGroupId());
        context.setArtifactId(model.getArtifactId());
        context.setBasePackage(model.getGroupId() + "." + model.getArtifactId());
        
        List<String> dependencies = new ArrayList<>();
        for (Dependency dep : model.getDependencies()) {
            dependencies.add(dep.getGroupId() + ":" + dep.getArtifactId() + ":" + dep.getVersion());
        }
        context.setDependencies(dependencies);
        
        File srcDir = new File(pomFile.getParentFile(), "src/main/java");
        if (srcDir.exists()) {
            context.setSrcRootPath(srcDir.getAbsolutePath());
            List<String> packages = parsePackages(srcDir);
            context.setPackages(packages);
        }
        return context;
    }

    private List<String> parsePackages(File srcDir) {
        List<String> packages = new ArrayList<>();
        File[] files = srcDir.listFiles();
        if (files == null) return packages;
        for (File file : files) {
            if (file.isDirectory()) {
                String packageName = file.getAbsolutePath().replace(srcDir.getAbsolutePath(), "")
                        .replace(File.separator, ".");
                if (!packageName.isEmpty()) {
                    packages.add(packageName.substring(1));
                }
                packages.addAll(parsePackages(file));
            }
        }
        return packages;
    }

    public boolean validateJavaCode(String code) {
        try {
            CompilationUnit cu = StaticJavaParser.parse(code);
            List<Problem> problems = cu.getProblems();
            return problems.isEmpty();
        } catch (Exception e) {
            return false;
        }
    }
}

@Data
public class ProjectContext {
    private String groupId;
    private String artifactId;
    private String basePackage;
    private List<String> dependencies;
    private List<String> packages;
    private String srcRootPath;
    private String projectId;
}

前端：IDEA 内的对话窗口

插件的工程搭建就不赘述了，新建 IntelliJ Platform Plugin 工程，选好 SDK 版本，在 plugin.xml 里注册 Action 和 ToolWindow。

<idea-plugin>
    <id>com.ai.code.assistant</id>
    <name>AI Code Assistant</name>
    <version>1.0</version>
    <vendor email="[email protected]">Your Name</vendor>
    <description>基于 Spring AI 的 Java 代码助手，支持上下文感知的代码生成与优化</description>
    
    <actions>
        <action text="AI Code Assistant" description="Open AI Code Assistant Dialog">
            <add-to-group group-id="EditorPopupMenu" anchor="first"/>
            <keyboard-shortcut keymap="$default" first-keystroke="ctrl alt A"/>
        </action>
    </actions>
    
    <extensions defaultExtensionNs="com.intellij">
        <toolWindow anchor="right" factoryClass="com.ai.code.assistant.window.AiToolWindowFactory"/>
    </extensions>
</idea-plugin>

对话窗口用 Swing 写，一个输入框、一个展示生成代码的文本区、生成和插入两个按钮。点击生成时，从当前编辑器拿到包名和 pom 路径，打包成请求发给后端。

public class AiCodeDialog extends JDialog {
    private JTextArea inputArea;
    private JTextPane resultArea;
    private JButton generateBtn;
    private JButton insertBtn;
    private Project currentProject;

    public AiCodeDialog(Project project) {
        super(WindowManager.getInstance().getFrame(project), "AI Code Assistant", Dialog.ModalityType.MODELESS);
        this.currentProject = project;
        initUI();
        setSize(800, 600);
        setLocationRelativeTo(null);
    }

    private void initUI() {
        inputArea = new JTextArea(5, 50);
        inputArea.setPlaceholder("请输入代码生成需求，例如：生成用户管理的 Controller+Service+Mapper");
        JScrollPane inputScroll = new JScrollPane(inputArea);

        resultArea = new JTextPane();
        resultArea.setContentType("text/java");
        JScrollPane resultScroll = new JScrollPane(resultArea);

        generateBtn = new JButton("生成代码");
        insertBtn = new JButton("插入到编辑器");
        insertBtn.setEnabled(false);

        JPanel panel = new JPanel(new BorderLayout());
        JPanel topPanel = new JPanel(new BorderLayout());
        topPanel.add(new JLabel("需求描述："), BorderLayout.NORTH);
        topPanel.add(inputScroll, BorderLayout.CENTER);

        JPanel btnPanel = new JPanel();
        btnPanel.add(generateBtn);
        btnPanel.add(insertBtn);

        panel.add(topPanel, BorderLayout.NORTH);
        panel.add(resultScroll, BorderLayout.CENTER);
        panel.add(btnPanel, BorderLayout.SOUTH);
        add(panel);
    }

    private void generateCode() {
        ProjectContext context = collectProjectContext();
        CodeGenerateRequest request = new CodeGenerateRequest();
        request.setRequirement(inputArea.getText());
        request.setProjectContext(context);

        OkHttpClient client = new OkHttpClient();
        // ... 发送请求并处理响应
        // String generatedCode = response.body().string();
        
        resultArea.setText(generatedCode);
        insertBtn.setEnabled(true);
    }

    private ProjectContext collectProjectContext() {
        ProjectContext context = new ProjectContext();
        String projectPath = currentProject.getBasePath();
        VirtualFile pomFile = currentProject.getBaseDir().findChild("pom.xml");
        if (pomFile != null) {
            context.setPomPath(pomFile.getPath());
        }
        Editor editor = FileEditorManager.getInstance(currentProject).getSelectedTextEditor();
        if (editor != null) {
            PsiFile psiFile = PsiDocumentManager.getInstance(currentProject).getPsiFile(editor.getDocument());
            if (psiFile instanceof PsiJavaFile) {
                PsiJavaFile javaFile = (PsiJavaFile) psiFile;
                context.setCurrentPackage(javaFile.getPackageName());
            }
        }
        context.setProjectId(currentProject.getName());
        return context;
    }

    private void insertCodeToEditor() {
        Editor editor = FileEditorManager.getInstance(currentProject).getSelectedTextEditor();
        if (editor == null) return;
        Document document = editor.getDocument();
        SelectionModel selectionModel = editor.getSelectionModel();
        int start = selectionModel.getSelectionStart();
        int end = selectionModel.getSelectionEnd();

        WriteCommandAction.runWriteCommandAction(currentProject, () -> {
            document.replaceString(start, end, resultArea.getText());
        });
        selectionModel.removeSelection();
        editor.getCaretModel().moveToOffset(start + resultArea.getText().length());
    }
}

代码生成：从一句话到完整代码

Prompt 工程

提示词的质量直接决定生成代码能不能用。我们在系统 prompt 里明确告诉 AI 项目的包名、已有的包、依赖列表，并要求只返回代码，不要解释。

@Service
public class PromptEngineeringService {
    
    public Prompt buildGeneratePrompt(String requirement, ProjectContext context) {
        String systemPrompt = "你是一位资深 Java 后端开发工程师，精通 Spring Boot、MyBatis、MySQL。\n" +
                "请根据以下需求和项目上下文，生成符合规范的 Java 代码：\n" +
                "1. 包结构必须符合项目基础包：%s\n" +
                "2. 代码必须兼容项目依赖版本，优先使用项目已引入的依赖\n" +
                "3. 生成完整的 Controller+Service+Mapper 层，包含必要的注释、异常处理\n" +
                "4. 代码风格符合阿里巴巴 Java 开发手册\n" +
                "5. 只返回代码，不返回多余解释\n" +
                "项目上下文：\n" +
                "- 基础包名：%s\n" +
                "- 已存在的包：%s\n" +
                "- 项目依赖：%s";
        
        String formattedSystemPrompt = String.format(systemPrompt,
                context.getBasePackage(),
                context.getBasePackage(),
                String.join(",", context.getPackages()),
                String.join(",", context.getDependencies()));

        String userPrompt = "需求：" + requirement;
        return new Prompt(List.of(
                new SystemMessage(formattedSystemPrompt),
                new UserMessage(userPrompt)
        ));
    }
}

生成接口

后端收到请求后，优先从缓存（Redis）拿项目上下文，没有则解析 pom 并存入。生成代码后用 JavaParser 做语法校验，不通过就让 AI 修复一次。

@RestController
@RequestMapping("/api/code")
public class CodeGenerateController {
    
    @Autowired
    private OpenAiChatClient openAiChatClient;
    @Autowired
    private PromptEngineeringService promptService;
    @Autowired
    private ProjectContextParser contextParser;
    @Autowired
    private ProjectContextRepository contextRepository;

    @PostMapping("/generate")
    public Result<String> generateCode(@RequestBody CodeGenerateRequest request) {
        try {
            ProjectContext context;
            if (request.getProjectContext().getProjectId() != null) {
                context = contextRepository.findByProjectId(request.getProjectContext().getProjectId());
                if (context == null) {
                    context = contextParser.parseMavenProject(request.getProjectContext().getPomPath());
                    context.setProjectId(request.getProjectContext().getProjectId());
                    contextRepository.save(context);
                }
            } else {
                context = request.getProjectContext();
            }

            Prompt prompt = promptService.buildGeneratePrompt(request.getRequirement(), context);
            AiResponse response = openAiChatClient.generate(prompt);
            String generatedCode = response.getGeneration().getText();

            boolean isValid = contextParser.validateJavaCode(generatedCode);
            if (!isValid) {
                generatedCode = regenerateCode(prompt);
            }
            return Result.success(generatedCode);
        } catch (Exception e) {
            log.error("生成代码失败", e);
            return Result.error("生成代码失败：" + e.getMessage());
        }
    }

    private String regenerateCode(Prompt prompt) {
        List<Message> messages = Stream.concat(
                prompt.getMessages().stream(),
                Stream.of(new UserMessage("以上代码存在语法错误，请修复后重新生成，只返回修复后的代码"))
        ).collect(Collectors.toList());
        Prompt newPrompt = new Prompt(messages);
        return openAiChatClient.generate(newPrompt).getGeneration().getText();
    }
}

输入'生成用户管理的 Controller+Service+Mapper，包含查询、新增、修改、删除接口'，会得到下面这样的 Controller：

package com.example.demo.controller;
import com.example.demo.entity.User;
import com.example.demo.service.UserService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import java.util.List;

@RestController
@RequestMapping("/api/user")
public class UserController {
    @Autowired
    private UserService userService;

    @GetMapping
    public List<User> listAll() {
        return userService.listAll();
    }

    @GetMapping("/{id}")
    public User getById(@PathVariable Long id) {
        return userService.getById(id);
    }

    @PostMapping
    public boolean save(@RequestBody User user) {
        return userService.save(user);
    }

    @PutMapping("/{id}")
    public boolean update(@PathVariable Long id, @RequestBody User user) {
        user.setId(id);
        return userService.update(user);
    }

    @DeleteMapping("/{id}")
    public boolean delete(@PathVariable Long id) {
        return userService.delete(id);
    }
}

代码优化：不只是修语法

优化功能的 Prompt 会告诉 AI 专注于某个优化类型，比如性能，并要求输出优化后的代码和说明。

public Prompt buildOptimizePrompt(String code, String optimizeType) {
    String systemPrompt = "你是一位资深 Java 性能优化工程师，精通 Java 语法、性能调优。\n" +
            "请根据指定类型优化以下代码：\n" +
            "优化类型：%s\n" +
            "优化规则：\n" +
            "1. 修复语法错误，保证代码可编译\n" +
            "2. 性能优化需给出具体的优化点（如循环优化、SQL 优化、集合使用优化）\n" +
            "3. 保留原有业务逻辑，只优化语法和性能\n" +
            "4. 输出优化后的代码 + 优化说明（分开展示）";
    
    String formattedPrompt = String.format(systemPrompt, optimizeType);
    return new Prompt(List.of(
            new SystemMessage(formattedPrompt),
            new UserMessage("需要优化的代码：\n" + code)
    ));
}

接口实现解析 AI 返回的'===优化后代码==='和'===优化说明==='分隔内容：

@PostMapping("/optimize")
public Result<CodeOptimizeResponse> optimizeCode(@RequestBody CodeOptimizeRequest request) {
    try {
        Prompt prompt = promptService.buildOptimizePrompt(request.getCode(), request.getOptimizeType());
        AiResponse response = openAiChatClient.generate(prompt);
        String result = response.getGeneration().getText();
        CodeOptimizeResponse responseVO = parseOptimizeResult(result);
        return Result.success(responseVO);
    } catch (Exception e) {
        log.error("优化代码失败", e);
        return Result.error("优化代码失败：" + e.getMessage());
    }
}

private CodeOptimizeResponse parseOptimizeResult(String result) {
    CodeOptimizeResponse response = new CodeOptimizeResponse();
    String[] parts = result.split("===优化说明===");
    if (parts.length >= 1) {
        response.setOptimizedCode(parts[0].replace("===优化后代码===", "").trim());
    }
    if (parts.length >= 2) {
        response.setOptimizeDesc(parts[1].trim());
    }
    return response;
}

比如把循环里查库的代码：

public List<User> listUsers(List<Long> ids) {
    List<User> users = new ArrayList<>();
    for (Long id : ids) {
        User user = userMapper.getById(id);
        users.add(user);
    }
    return users;
}

优化后变成批量查询，并加了空判断：

public List<User> listUsers(List<Long> ids) {
    if (CollectionUtils.isEmpty(ids)) {
        return Collections.emptyList();
    }
    return userMapper.listByIds(ids);
}

上下文感知补全

补全功能在插件端监听编辑器输入，取光标前 100 个字符作为代码片段，同时抓出当前包名和 import 列表，发给后端。后端 prompt 要求生成简洁的补全建议，每行一个。

@Service
public class CodeCompletionService {
    
    public List<String> completeCode(CodeCompletionRequest request) {
        String systemPrompt = "请根据当前 Java 文件的上下文，生成代码补全建议：\n" +
                "1. 补全建议必须符合当前包结构：%s\n" +
                "2. 优先使用已导入的类：%s\n" +
                "3. 补全建议简洁，每条不超过 50 个字符\n" +
                "4. 只返回补全建议列表，每行一个";
        
        String formattedPrompt = String.format(systemPrompt,
                request.getCurrentPackage(),
                String.join(",", request.getImportedClasses()));
        
        String userPrompt = "需要补全的代码片段：\n" + request.getCodeSnippet();
        Prompt prompt = new Prompt(List.of(
                new SystemMessage(formattedPrompt),
                new UserMessage(userPrompt)
        ));

        AiResponse response = openAiChatClient.generate(prompt);
        String result = response.getGeneration().getText();
        
        return Arrays.stream(result.split("\n"))
                .map(String::trim)
                .filter(s -> !s.isEmpty())
                .collect(Collectors.toList());
    }
}

插件端异步调用补全 API，然后把建议显示出来（具体显示方式略）。

public class CodeCompletionListener extends TypedActionHandlerBase {
    @Override
    public void execute(@NotNull Editor editor, char c, @NotNull DataContext dataContext) {
        CaretModel caretModel = editor.getCaretModel();
        int offset = caretModel.getOffset();
        Document document = editor.getDocument();
        String codeSnippet = document.getText(new TextRange(Math.max(0, offset - 100), offset));

        Project project = CommonDataKeys.PROJECT.getData(dataContext);
        PsiFile psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document);
        CodeCompletionRequest request = new CodeCompletionRequest();
        request.setCodeSnippet(codeSnippet);
        
        if (psiFile instanceof PsiJavaFile) {
            PsiJavaFile javaFile = (PsiJavaFile) psiFile;
            request.setCurrentPackage(javaFile.getPackageName());
            List<String> importedClasses = javaFile.getImportList().getAllImports().stream()
                    .map(ImportStatement::getQualifiedName)
                    .collect(Collectors.toList());
            request.setImportedClasses(importedClasses);
        }

        CompletableFuture.runAsync(() -> {
            List<String> completions = callCompletionApi(request);
            showCompletionSuggestions(editor, completions);
        });
    }
}

打包部署与私有化

插件打包用 Gradle 的 org.jetbrains.intellij 插件，配置一下 SDK 版本和输出目录：

plugins {
    id 'java'
    id 'org.jetbrains.intellij' version '1.17.3'
}

intellij {
    version = '2023.2'
    type = 'IC'
    plugins = ['java']
}

sourceCompatibility = 17
targetCompatibility = 17

tasks.buildPlugin {
    archiveBaseName = 'ai-code-assistant'
    archiveVersion = '1.0.0'
    destinationDirectory = file("$projectDir/dist")
}

然后 ./gradlew buildPlugin，生成的 zip 上传到私有 Nexus 或插件仓库。在 IDEA 的插件设置里添加私有仓库地址，就能搜索安装。

后端 Spring Boot 服务打成 jar，nohup java -jar ... & 启动，前面挂个 Nginx 反代就行。

实际开发中遇到的坑

把踩过的几个典型问题列出来，供参考：

• 插件启动时取不到项目上下文：因为插件初始化太早，项目还没加载完。改成在 projectOpened 事件里触发上下文采集就解决了。
• AI 生成的代码包名错误：Prompt 里虽然指定了基础包，但有时 AI 还是会用默认的 com.example。我们在 prompt 里反复强调，同时后端做正则校验，强制替换包名。
• 生成响应太慢：首次请求时要解析 pom 和源码目录，加上 AI 调用，经常超过 5 秒。后来把项目上下文缓存到 Redis，并改成异步：生成接口立刻返回任务 ID，前端轮询结果，体验好多了。
• JavaParser 校验误判：JavaParser 版本和 IDEA 的 PSI 解析器有细微差异，偶尔会把好代码判错。最终直接改用 IDEA 内置的 PSI API 做校验，保持一致性。
• 插入代码后格式错乱：因为换行符和缩进不一致。插入前用 CodeStyleManager.getInstance(project).reformat(psiElement) 统一格式化一下就好。

后续方向

目前这个助手已经能在项目中用起来，后面会往这几个方向拓展：

• 接入国产模型私有化部署，减少数据外泄顾虑。
• 把数据库表结构、接口文档作为知识库喂给模型，生成更精准的代码。
• 支持根据表结构一键生成整个模块。
• 团队协作功能，共享 Prompt 模板和代码片段。
• 适配 VS Code 等其他 IDE。

整个开发过程并不复杂，关键是把项目的上下文信息准确塞进 AI 的提示词里。如果你也有类似需求，可以参考这套方案动手试试。