Rivet improves code generation in existing software repositories
Rivet: Utility-Preserving Context Control for Repository-Level Coding Agents
Software Engineering
Summary
It can be hard for AI coding assistants to add new code to a big software project because they must understand important information like where the code goes and how it connects to other parts. The authors identify that simply picking parts of the project that look relevant can miss this critical information. They created Rivet, a method that carefully picks and summarizes the right clues from the project so the AI can write better code within its limited memory. Tests show Rivet helps AI write working code more often and reduces computing costs.
What this means in practice
- •For software development teams: Improve AI coding assistants’ ability to add correct code inside large existing projects by better selecting relevant context.
- •For devops teams: Reduce computing costs for AI-assisted code generation in repositories by efficiently managing context selection and code summarization.
Authors
Ruwei Pan, Linhao Wu, Yakun Zhang, Zixiong Yang, Lu Zhang, Hongyu Zhang
Abstract
Repository-level code generation asks an agent to implement a task requirement inside an existing repository. To produce code that works in the repository, the agent needs repository evidence, such as target locations, helper definitions, caller usages, and configuration constraints, which determine how the code should be located, integrated, and executed. Because the context budget cannot hold the entire repository, existing retrieval, compression, and structure-aware methods select context by task similarity, context length, or structural proximity, but these criteria do not directly measure whether the selected context contains this evidence. As a result, the selected context can appear relevant to the task yet miss the evidence that the implementation depends on. We call this failure mode context utility dilution. In this paper, we propose Rivet, a context control framework that preserves repository evidence for coding agents under a fixed context budget. Rivet, first uses Signal Grounding to select task-related code anchors by fusing lexical, semantic, and structural signals. Utility Linking then expands these anchors through dependency, usage, and configuration relations to collect repository evidence. Finally, Budget Shaping packs the evidence by priority and, when the budget is tight, reduces snippets to signatures or summaries before dropping them. We evaluate Rivet, on two executable repository-level benchmarks, DevEval and RepoExec, with four backbone models, two context budgets (4K and 8K tokens), and eight baselines. Rivet, achieves the highest Pass@1 in all 16 settings and improves over the strongest baseline by 3.73--5.15 percentage points on average. Moreover, GPT-5-mini with Rivet, outperforms GPT-5 without Rivet,on both benchmarks while reducing inference cost by 73.3%--84.5%.