Security category aware retrieval improves code generation defense
CAT-RAG: Improving Secure Code Generation with Category-Aware Agentic Retrieval
Cryptography and SecuritySoftware Engineering
Summary
Writing code with fewer security problems is still a challenge even with advanced AI code helpers. The authors designed a new method called CAT-RAG that helps AI models pick the right security advice by grouping vulnerabilities into simpler categories instead of many detailed weaknesses. This way, the AI focuses better and produces safer code using less extra information. In tests, this method nearly doubled improvements in secure code compared to prior techniques and reduced extra content sent to the AI by about 40%.
What this means in practice
- •For software development teams: Produce more secure code automatically by integrating smarter vulnerability category guidance into their code generation tools.
- •For cybersecurity analysts: Enhance automated code security reviews by categorizing and consolidating security knowledge to better guide vulnerability detection and fixing.
Authors
Mahdieh Abazar, Lorenzo De Carli, Benjamin Tan
Abstract
Despite rapid progress in LLM-based code generation, security vulnerabilities remain widespread, and retrieval-augmented code generation (RACG) can further amplify risk when retrieved knowledge is verbose or vulnerable to poisoning. Existing secure RACG approaches inject security knowledge at a fine-grained vulnerability level, causing prompt overhead and diluted model attention when a query implicates multiple related weaknesses. We propose CAT-RAG, an agentic security category-aware retrieval-augmented code generation framework. CAT-RAG performs adaptive category selection via an LLM as a categorizer agent and operates either as a standalone security layer or as a drop-in add-on to RACG pipelines, requiring no modification to the underlying model. We replace CWE-level retrieval with compact, language-specific security categories. These categories provide a higher-level abstraction over individual CWE weaknesses, organizing them by the security concern they share, such as Input Validation and Injection Attacks, rather than merely grouping related entries. Each category consolidates best practices and secure code examples into a single, prompt-efficient representation. We perform a comprehensive evaluation across CyberSecEval, CodeGuard+, and HumanEval. On CyberSecEval, CAT-RAG nearly doubles the secure code generation improvement over the state of the art (23.57% vs. 12.11%) while reducing token overhead by 41.71%.