COSTA: Code Watermarking via Statement-Level ASTs for Large Language Models

Large language models exhibit strong code generation capabilities, accelerating software development and widespread adoption of AI coding assistants. As generated code becomes harder to distinguish from human-written code, legal, ethical, and security concerns motivate watermarking for provenance tracing. However, existing generation-time code watermarking methods rely mainly on token-level biases that can compromise functional correctness, while sentence-level text watermarking remains difficult to transfer directly from natural language to constrained code space. In this paper, we introduce COSTA, the first generation-time, statement-level code watermarking framework that encodes provenance in the abstract syntax trees (ASTs) of complete statements. Specifically, a secret key and a locality-sensitive hash of the preceding statement jointly determine the target parity of a selected AST node count. The generator requires no model retraining and embeds watermark signal through three stages: sample, repair, and bridge. The detector tests whether keyed parity matches occur more often than chance, requiring neither the original prompt nor access to the generator. Extensive experiments across models, benchmarks, programming languages, and code transformations demonstrate that COSTA combines strong detectability and robust watermark retention while largely preserving the quality of generated code.