Reflection-Driven Control for Trustworthy Code Agents

🕒 Published (v1): 2025-12-22 00:27 UTC · Source: Arxiv · Venue: AAAI 2026 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Reflection-Driven Control (Reflex) is a plug-and-play agent module that promotes self-reflection from a post-hoc patch into a first-class control loop embedded within the agent's reasoning pipeline for secure code generation. It combines a lightweight binary safety classifier, a multi-turn reflective prompt engine, and a hierarchical dynamic/static RAG memory to detect and repair security vulnerabilities before code is committed. Across eight CWE vulnerability categories and four LLMs, it improves Security Rate by 2.9–11.2 percentage points with negligible token and latency overhead.

Problem

Contemporary LLM code agents produce unsafe outputs in uncontrolled ways: even strong base models hallucinate or emit vulnerable code, and agentic tool-use pipelines amplify risks through jailbreaks, prompt injection, and toolchain hijacking. Existing safety mechanisms are reactive or post-hoc (external auditors, fine-tuning, input-layer reminders) and lack auditable, real-time control integrated into the agent's own reasoning loop.

Method

The Reflex module inserts three components into any existing Plan–Execute–Verify agentic architecture:

1. Lightweight Self-Checker: A single LLM call classifies input code as SAFE/UNSAFE via a concise safety-review prompt. SAFE samples are written directly to memory; UNSAFE samples trigger the full reflection pipeline, avoiding unnecessary computation on safe inputs.

2. Reflective Prompt Engine: Constructs a structured 3-round chain-of-thought dialogue — (Round 1) vulnerability identification, (Round 2) root-cause and mitigation, (Round 3) secure code revision — making the full reasoning chain explicit and auditable.

3. Reflective Memory Repository: A hierarchical two-tier RAG store: Dynamic Memory (ChromaDB vector store) accumulates verified repair cases at runtime; Static Memory holds predefined secure coding standards and CVE databases. Retrieval prioritizes dynamic memory (cosine similarity ≥ 0.70, at least k_min results); it falls back to static memory when dynamic recall is insufficient. Verified outputs are written back, forming a self-reinforcing knowledge loop. The module requires no base-model fine-tuning.

Key Contributions

• A closed-loop, reflection-as-control-circuit framework spanning planning, execution, and verification stages with an auditable evidence trail.
• A practical instantiation for secure code generation combining lightweight self-check routing, dynamic RAG-based reflective prompting, and tool-gated verification (compiler + CodeQL).
• Empirical evaluation across 8 CWE categories × 4 LLMs showing consistent security gains with cost and latency breakdown demonstrating near-zero marginal overhead.
• Analysis showing knowledge saturation after ~4 rounds (≈100 samples) and that a single reflection round captures ~90% of available repair patterns.

Results

• Security Rate improvements (Base → Base+Reflex): GPT-3.5-turbo: 93.7→96.6% (+2.9pp); GPT-4o: 85.7→95.0% (+9.3pp); Qwen3-Coder-Plus: 83.7→94.9% (+11.2pp); Gemini-2.5-Pro: 88.0→97.1% (+9.1pp).
• Pass Rate: Mixed — GPT-3.5-turbo +4.4pp, GPT-4o −0.3pp, Qwen3-Coder-Plus −6.6pp, Gemini-2.5-Pro +3.5pp; functional correctness is partially traded off for security in weaker models.
• Dynamic RAG convergence: Average retrieval similarity improved from 0.850 (run 1) to 0.980 (run 5); retrieval success rate 85%→100%; fallback rate 15%→0% by run 4.
• Repair accuracy by similarity: Docs with cosine sim ≥0.95 achieve 100% FixAcc; 0.85–0.95 → 98.6%; 0.70–0.85 → 93.8%; <0.70 → 75.0%.
• Cost: Average $5.37×10⁻⁴ per scenario; total 44,762 tokens and ~250 API calls over 125 scenarios.
• Latency: 28.8 s/scenario average; LLM inference dominates (84.4%); dynamic RAG retrieval only 0.8 s (2.8%).

Limitations

• Pass Rate degrades on some models (Qwen3-Coder-Plus −6.6pp), indicating a security–functionality trade-off that is not consistently resolved.
• Evaluated only on a single benchmark (He & Vechev 2023 CWE dataset, 8 vulnerability types, C/C++ and Python); generalization to other vulnerability classes or languages is unverified.
• The 3-round reflection dialogue structure is fixed; optimal depth may vary by vulnerability type.
• Batch experiments ran ~2× longer than theoretical estimates due to data persistence and state management overhead, suggesting production deployment requires additional engineering.
• Dynamic memory quality and diversity depend on seed inputs; cold-start behavior with zero prior cases is not characterized in depth.
• No red-team evaluation of whether adversarial inputs can corrupt or poison the dynamic memory repository.

Relevance to Agentic AI / LLM Agents

This work directly addresses the agentic safety gap where tool-using, multi-step LLM agents can produce harmful artifacts without intrinsic self-supervision — a central concern in the TRiSM and OWASP-LLM threat models for agentic systems. The Reflex module's plug-and-play design is architecturally significant: it shows that safety control can be injected as a composable submodule rather than requiring model retraining or a dedicated external monitor, offering a practical template for trust-by-construction in agent orchestration. The hierarchical dynamic RAG memory that accumulates verifiable repair cases over time is a concrete instance of agentic continual learning with auditability. For researchers tracking agentic AI, this paper sits at the intersection of agent safety, RAG-augmented self-correction, and code agent evaluation — providing both a deployable mechanism and an ablation-validated understanding of reflection depth dynamics.