Gallery inside!
Research

AI Agent Safety: Test the Memory, Tools and Actions Around the Model

AI agent safety extends beyond model refusals. Use research on prompt injection, memory and tools to test permissions, actions and useful task completion.

7

Select a figure to open it at full size.

An assistant that only answers questions has a different risk profile from one that reads documents, remembers instructions and changes records. Even if both use the same model, the second system has more ways to turn a mistaken interpretation into an unwanted action.

Safety at Scale surveys attacks, defenses and benchmarks across large models and agents. Its most useful lesson for product teams is to evaluate the surrounding system, not only the model’s willingness to refuse a harmful prompt.

A tool connection, a memory store or another agent can introduce information that the model treats as authoritative. The engineering question is where that information is allowed to influence decisions—and which actions remain constrained outside the model.

What the Research Actually Shows

Xingjun Ma and colleagues organize 574 technical papers identified through keyword searches and manual filtering. The survey covers vision models, language models, multimodal systems, diffusion models and agents. Within those categories it distinguishes threats such as adversarial inputs, jailbreaks, prompt injection, poisoning and data extraction.

This is a structured review of published research. The authors do not rerun every method on a common system, and the paper does not supply a universal safety score or a measured business return from adopting a particular defense.

That distinction matters because the underlying studies make different assumptions. Some require access to model internals; others only interact through an API. Some assess generated text; others assess tool execution. Results from those settings cannot be treated as interchangeable product rankings.

The survey’s agent taxonomy is especially useful when an assistant gains new capabilities.

Original agent branch of the survey taxonomy listing indirect prompt injection, memory, tools, visual agents, multi-agent systems, embodied agents and benchmarks.
Agent branch excerpt from original Figure 3, Ma and colleagues. The separate rows identify places a model-only safety test can miss. This is a classification of research, not a ranking of threat frequency. Full original figure.

The Core Insight: Untrusted Content Can Become an Instruction

A jailbreak typically tries to make a model violate its behavioral restrictions. Indirect prompt injection reaches the model through something it is supposed to read: a retrieved page, a document or a tool result.

The problem is the transition from data to authority. The assistant needs the document’s information to complete the user’s task, but the document may also contain text directing the assistant to do something else. If those roles are not distinguished, useful retrieval becomes a path for task redirection.

The survey describes several responses. Instruction hierarchy gives different message sources different authority. Provenance-marking methods make the origin of external content clearer. Isolation patterns restrict which components can see sensitive information or call particular tools. These approaches address different parts of the problem; none follows automatically from adding a stronger refusal prompt.

Consider a customer-support assistant summarizing an uploaded file. A malicious instruction inside that file should remain content to analyze. It should not become permission to change the customer’s account or send information elsewhere. This example shows the product boundary: the user authorized a summary, not every action mentioned in the material being summarized.

The indirect-injection discussion explains why an apparently benign user request does not make every downstream input trustworthy.

Memory Makes the Problem Persist

The survey separates immediate context from longer-lived memory. Both can affect later behavior, but persistent memory creates an additional path: a problematic interaction can become stored guidance that is retrieved on a future task.

Its memory-poisoning examples include false interaction records and misleading material placed in retrieval collections. The mechanism does not always require changing model weights. The attacker influences what the system remembers or retrieves, and the model then reasons from that material as if it were useful history.

For a product team, the important sequence is input, storage, retrieval, reuse. Checking a response once does not establish that the same information is safe to save as an enduring preference or operating instruction.

A practical application is to give memory records an origin and a scope. A confirmed user preference, a quoted document and an inferred summary should not acquire the same authority merely because they share a database. When a record is corrected or removed, the team should also check whether derived summaries continue to carry it forward.

The survey discusses defenses based on access control, separating and aggregating sources, and resolving conflicts between retrieved and internal information. These are research directions with different assumptions—not evidence that any one memory store becomes trustworthy by default. The memory sections develop this distinction.

Tool Access Changes the Consequences

A misleading answer can be corrected before anyone acts on it. An incorrect tool call may already have changed a record, disclosed data or triggered a workflow.

The survey treats tool descriptions, selection, inputs, execution and outputs as potential points of failure. It also discusses MCP-related research, where information describing a tool or server can influence which capability an agent selects. A standard connection protocol does not itself establish that the connected capability is appropriate for the user’s task.

The resulting design question is concrete: what can this model request, and what will the application actually allow? Restricting available actions, checking arguments and enforcing authorization in application code creates boundaries that do not depend entirely on the model interpreting every sentence correctly.

Adding a second agent is not automatically a defense. The multi-agent section describes how misleading instructions can travel when agents quote or reuse each other’s messages. Collaboration can provide additional scrutiny, but it can also distribute the same mistake. Each handoff needs to preserve the distinction between evidence, a proposal and an authorized action.

Build an Evaluation Around the Failure You Care About

The authors distinguish lightweight simulated evaluations from tests involving tool or environment interaction. The first can cover many cases cheaply. The second can reveal whether a supposedly safe response still leads to an unsafe action.

For the support example, test more than whether the final answer contains a refusal. Inspect the action trace: was a write requested, was it allowed, what data was passed, and did the original summary task still succeed?

A useful evaluation set contains legitimate tasks, malicious external content within otherwise legitimate tasks, and ambiguous cases requiring clarification. Keep the ordinary tasks: a defense that blocks every document may prevent misuse while making the product unusable.

The evaluation section also argues that attack success rate alone is insufficient. Failures differ in severity, recoverability and effect. Record those outcomes alongside task completion, unnecessary refusals, latency and review burden. Keep fresh cases outside routine tuning so success does not merely reflect familiarity with a fixed test set.

Implementation Frameworks

PyRIT is an open-source risk-identification framework for generative AI testing. It is appropriate when a team needs a repeatable evaluation process across model or application changes. It is an evaluation tool, not a runtime authorization boundary.

For a first pass, use a sandbox version of the actual workflow with synthetic documents and non-production accounts. Define expected permitted actions for each task. Run the current system and proposed controls against the same cases, then inspect both responses and tool traces. A small existing test harness can be sufficient if it captures these outcomes reliably; installing a larger framework is not the result being sought.

Application permissions and explicit action validation serve a different role. Keep the model’s proposal separate from the code that determines whether an action is allowed. Start with the highest-consequence capability: define its permitted scope, validate its arguments and record the decision. A model-based guard may add useful screening, but it should not silently become the only source of authority.

For selecting research evaluations, the survey’s maintained reference collection helps locate methods by threat and model type. Read a chosen benchmark’s assumptions before adopting its score as a release criterion. Matching the application’s failure mode matters more than using the longest checklist.

TechClarity’s View

Expanding an assistant’s capabilities should trigger a new safety evaluation. Memory, tools and agent collaboration alter the system even when the underlying model is unchanged.

The survey is valuable because it makes those boundaries visible. It cannot certify a deployment, and the volume of published attack research does not tell a team how often each threat will occur in its product. Use the taxonomy to find missing tests, then judge controls by the consequences they prevent and the legitimate work they preserve.

Our guardrails analysis examines the narrower tradeoff between screening and usability. This broader system view explains why that screening is only one part of the design.

Original Research

Safety at Scale: A Comprehensive Survey of Large Model and Agent Safety, Xingjun Ma and colleagues. This article uses arXiv version 6, submitted 14 April 2026.

Related research

Tags:
Author
TechClarity Analyst Team
September 27, 2026