💻 computer science

MAGE: Safeguarding LLM Agents against Long-Horizon Threats via Shadow Memory

The paper introduces MAGE, a novel defensive framework that leverages a dedicated "shadow memory" to proactively distill safety-critical context and detect long-horizon threats in LLM agents, achieving high detection accuracy with negligible overhead.

Original authors: Yuhui Wang, Tanqiu Jiang, Jiacheng Liang, Charles Fleming, Ting Wang

Published 2026-05-06

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuhui Wang, Tanqiu Jiang, Jiacheng Liang, Charles Fleming, Ting Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Slow Poison" Attack

Imagine you hire a very smart, helpful robot assistant (an LLM Agent) to do your daily tasks, like checking your email, managing files, or sending messages.

Usually, we worry about someone yelling a bad command at the robot all at once (like "Delete everything!"). But this paper talks about a sneakier, more dangerous problem called Long-Horizon Threats.

Think of this like slow poison.

Step 1: The attacker asks the robot to find a specific file. (Totally normal).
Step 2: The attacker asks the robot to look up a friend's email address. (Totally normal).
Step 3: The attacker asks the robot to email that file to that friend. (Totally normal).

Individually, every step looks innocent. But when you put them together, the robot just accidentally leaked your secret company strategy to a competitor. The robot gets confused because it forgets the context of Step 1 by the time it reaches Step 3. It's like a movie where the villain tricks the hero into doing small, harmless things that eventually lead to a disaster, but the hero only remembers the current scene, not the whole plot.

The Solution: MAGE and the "Shadow Memory"

The authors created a defense system called MAGE. To understand how it works, imagine a high-security bank vault.

Normally, the bank teller (the Agent) looks at the customer's request and decides if it's okay. But if the customer has been whispering tricks to the teller over the last hour, the teller might get confused.

MAGE introduces a "Shadow Memory."
Think of this as a security guard sitting in a glass booth next to the teller.

The teller (the Agent) does the work.
The security guard (MAGE) doesn't do the work, but they watch everything that happens.
Crucially, the guard keeps a special, condensed notebook (the Shadow Memory).

Every time the teller does something, the guard writes a tiny, bullet-point summary in their notebook. They don't write down the boring details (like "the weather was nice"); they only write down safety-critical clues (like "User is asking for a confidential file" or "Recipient is an outsider").

Before the teller actually sends an email or deletes a file, they have to ask the guard: "Hey, based on my notebook of everything that happened today, is this safe?"

If the guard sees that Step 1 was "Find Secret File" and Step 2 was "Email to Stranger," they will slam the brakes and say, "NO! This is a trap!" even if the current request (Step 3) looks harmless on its own.

How MAGE is Different

Most previous security systems are like a keyword filter. They only look at the current sentence. If the sentence says "Send email," they let it pass. They miss the big picture.

MAGE is different because:

It remembers the whole story: It connects the dots between the first step and the last step.
It's efficient: Instead of reading the entire history of the conversation (which is huge and slow), the guard only reads their tiny, condensed notebook. This makes it fast and cheap.
It learns: The system was trained using a method called "Reinforcement Learning." Think of this like a video game where the guard gets points for catching bad guys and loses points if they accidentally stop a good guy. Over time, the guard gets really good at spotting the "slow poison" attacks without being too paranoid.

The Results: What Happened in the Tests?

The researchers tested MAGE against two types of bad guys:

The User: Someone trying to trick the robot by asking a series of innocent questions that lead to a bad outcome.
The Environment: Someone hiding bad instructions inside the data the robot reads (like a website the robot visits).

The findings were impressive:

Caught the bad guys: MAGE stopped almost all of these "slow poison" attacks. In one test, it reduced the success rate of attacks from 100% down to less than 10%.
Didn't stop the good guys: It didn't get in the way of normal work. The robot could still finish its tasks almost as well as it could without any security guard.
Caught them early: MAGE usually spotted the danger right at the beginning of the attack, giving the human operator time to step in before any damage was done.
Low cost: It didn't slow the robot down much or cost a lot of computer power.

The Bottom Line

The paper argues that as AI agents get smarter and do longer, more complex tasks, they become vulnerable to attacks that unfold over time. MAGE solves this by giving the agent a "second brain" (the Shadow Memory) that specifically watches for safety risks across the entire timeline of a task, ensuring that a series of small, innocent steps doesn't accidentally turn into a disaster.

Technical Summary: MAGE – Safeguarding LLM Agents against Long-Horizon Threats via Shadow Memory

1. Problem Statement

As Large Language Model (LLM) agents are deployed for complex, multi-turn tasks, they face a growing class of long-horizon threats. Unlike single-turn attacks (e.g., standard jailbreaks), these threats exploit extended interactions where malicious intent is latent within individual turns but emerges destructively only at the trajectory level.

The paper identifies three primary causes for this vulnerability:

Context Dilution: The tension between finite working memory and growing context windows causes security-critical information to be diluted over time.
Lack of Retrospective Inspection: Agents often lack the capacity to inspect the full trajectory to detect adversarial intent that accumulates gradually.
Training Distribution Shift: Safety alignment is typically trained on short contexts and fails to generalize to extended interactions.

Existing defenses (e.g., keyword filters, per-turn classifiers, safety prompting) degrade significantly under these long-horizon conditions. While some recent works attempt multi-turn alignment, they often fail when interaction lengths exceed training regimes or focus narrowly on jailbreaks, leaving broader threats like Sequential Tool-Attack Chaining (STAC) and Persistent Indirect Prompt Injection (PI2) unaddressed.

2. Methodology: The MAGE Framework

The authors propose MAGE (Memory As Guardrail Enforcement), a defensive framework inspired by the "shadow stack" abstraction in systems security. Instead of relying solely on the agent's working memory, MAGE introduces a dedicated, security-focused Shadow Memory that runs in parallel.

Core Architecture

MAGE consists of two primary components, both implemented as tool calls within the agent's workflow and trained via Reinforcement Learning (RL):

Memory Manager ( $M$ ): A module that distills and maintains the shadow memory. It processes the agent's execution trajectory (user instructions, actions, and observations) to extract and retain security-critical context while discarding irrelevant information. It tracks user intent and known attack patterns.
Judge ( $J$ ): A module that performs retrospective risk assessment. Before an action is executed, the Judge consults the current state of the Shadow Memory to assess the risk of the pending action. It outputs a decision (approve/reject) and a human-readable rationale.

Workflow

At each turn $t$ :

The agent proposes an action $a_t$ .
The Memory Manager updates the shadow memory $m_t$ based on the previous state $m_{t-1}$ and the current trajectory segment.
The Judge assesses the risk of $a_t$ conditioned on $m_t$ .
If the Judge rejects the action, execution is halted, and the rationale is reported. Otherwise, the action proceeds.

Optimization via Turn-wise RL

The framework employs Group Relative Policy Optimization (GRPO) to jointly train $M$ and $J$ . The training is unique in its turn-wise granularity:

Reward Design: The system assigns rewards at the turn level rather than the trajectory level.
- Judge Reward: Binary correctness based on ground-truth labels for the specific action.
- Memory Reward: A temporally discounted reward that propagates future judge outcomes backward to the memory updates that contributed to them. This prevents myopic memory updates and ensures the memory retains information critical for future detection.
- Penalties: Format and length penalties are applied to ensure structured outputs and compact memory representations.
Phase Separation: The loss function separates the generation phases of the Memory Manager and the Judge, ensuring that rewards are attributed correctly to the specific tokens generated by each component.

3. Key Contributions

Identification of the Core Challenge: The paper posits that defending against long-horizon threats requires maintaining security-critical context across the entire execution trajectory, a capability missing in standard agent designs.
Novel Framework (MAGE): MAGE is the first framework to leverage agentic memory as a security primitive. It introduces a shadow memory mechanism to distill context and perform retrospective, trajectory-level risk assessment.
Plug-and-Play Design: The framework integrates with existing LLM agents with minimal modification to the core inference pipeline, utilizing the agent's native tool-calling capabilities.
Empirical Validation: Extensive evaluation demonstrates MAGE's superiority over existing defenses in detection accuracy, early-stage detection, and utility preservation.

4. Experimental Results

The authors evaluated MAGE under two threat models: User-as-Adversary (STAC attacks) and Environment-as-Adversary (Indirect Prompt Injection, including PI2).

Evaluation I: User-as-Adversary (STAC)

Setup: Used GPT-5.1 as the agent backbone and Qwen3-4B for MAGE components.
Performance: MAGE reduced the Attack Success Rate (ASR) of STAC attacks from 100.0% (no defense) to 8.3%, while preserving 94.4% of benign utility (BU).
Efficiency: Introduced only 7.0K extra tokens per task, significantly lower than baseline variants.
Detection Latency: Detected the majority of attacks at or immediately adjacent to the actual malicious action (average delay of -0.06 turns), providing early warnings.
Ablation: Variants without shadow memory or without RL training showed significantly degraded performance (e.g., untrained variants achieved 0% ASR but also 0% BU due to indiscriminate rejection).

Evaluation II: Environment-as-Adversary (PI2)

Setup: Evaluated against persistent indirect prompt injection (PI2) and static injection attacks across GPT-4o-mini, Llama-3.3, and Qwen3.
Performance: MAGE achieved 0.0% ASR on PI2 attacks (where baselines like Spotlighting and Repeated-Prompt failed with >70% ASR) and maintained high benign utility (73.0% on Llama-3.3, 86.5% on Qwen3).
Comparison: MAGE outperformed baselines (PI-Detector, Spotlighting, MELON, IPIGuard) in the security-utility-overhead trade-off, achieving the lowest ASR and high BU with low token overhead.
Adaptive Attacks: Even under adaptive attacks where the adversary knows MAGE's architecture and attempts to poison the memory or override the judge, MAGE maintained robust performance (ASR remained low, e.g., 0.0% for PI2 adaptive attacks), whereas untrained variants or those without shadow memory failed.

5. Significance and Claims

The paper claims that MAGE establishes a new paradigm for defending LLM agents against long-horizon threats. Its significance lies in:

First-of-its-kind Approach: It is the first work to utilize agentic memory specifically as a defensive primitive for trajectory-level threat detection.
Balanced Trade-offs: It successfully addresses the triad of Security (low ASR), Utility (high BU), and Overhead (low token cost), a balance that existing defenses fail to achieve simultaneously.
Interpretability: By grounding detection decisions in the shadow memory, MAGE provides human-readable explanations, making the defense transparent to operators.
Generalizability: The framework is designed to be plug-and-play, requiring no modification to the agent's core inference pipeline and functioning effectively with smaller, local models for the defense components.

The authors conclude that while MAGE represents a significant step forward, future work should explore broader attack classes (e.g., memory poisoning, goal drifting), dynamic memory schemas, and extensions to multi-agent systems.