Skip to content

Latest commit

 

History

5 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

agentmemorytaintgap

Catch untrusted content being written into an AI agent's long-term memory as if it were a trusted fact.

Prompt injection usually gets pictured as a single bad turn: a poisoned web page or tool result sneaks into one prompt, the model says something wrong once, and the blast radius ends there. Agent memory breaks that assumption. A growing class of agent frameworks let the agent persist content — a tool result, a summary, a "fact" — into a long-term memory store that gets read back and fed into every future prompt as if it were the agent's own trusted conclusion. If the content that gets written was actually attacker-controlled (a malicious tool response, a crafted user message, a poisoned web page), and nothing strips or tags it first, the poison survives for the lifetime of that memory store — this is the "agent memory poisoning" risk that's an increasingly-discussed, distinct branch of prompt injection in 2026.

agentmemorytaintgap reads your source with Python's ast module — no imports, no execution — and flags the memory-write call sites where that can happen:

$ agentmemorytaintgap agent/

BLOCKER  AT001 agent/handler.py:3:4   Memory-write call stores content traced back to an untrusted origin (tool output / external fetch / raw user input) with no sanitization or provenance-tagging call in this function (agent memory poisoning risk).
        agent_memory.save_context({"input": user_query}, {"output": tool_result})

1 file(s) scanned · 1 blocker(s) · 0 warning(s)

Exit code 1 on a blocker, so it drops straight into pre-commit or CI.

What it flags

Rule Severity Fires when…
AT001 blocker a memory-write call's stored value traces (single-hop, same function) back to a provably untrusted origin — tool output, an HTTP/fetch response, or a raw user-input parameter — with no sanitize/tag call on it anywhere in that function.
AT002 warning a memory-write call's stored value origin could not be confidently traced either way (a bare parameter with no naming signal, or a longer assignment chain). Lower confidence — worth a human glance, not a confident blocker.

The LLM's own generated response (e.g. llm.invoke(...)) is treated as trusted and is never flagged — the concern here is specifically external content being stored as if it were the agent's own conclusion.

Recognized memory-write shapes

  • LangChain-shaped: .save_context(...), .chat_memory.add_message(...) / .add_user_message(...) / .add_ai_message(...).
  • Generic vector-store-as-memory: .add(...) / .upsert(...) called on a variable whose name contains memory, mem_store, or long_term — there is no single standard "agent memory" API the way there is for HTTP, so this is a naming heuristic, documented honestly in DETAILS.md.
  • Custom memory helpers: a call to remember(...), store_memory(...), or save_memory(...), as a free function or bound method.

Recognized untrusted origins

  • A tool call: .run(...) / .invoke(...) on a variable named like tool, or a call to a function decorated with a recognizable @tool decorator.
  • An external fetch: requests/httpx/aiohttp .get/.post/...(...), or a call to a function literally named fetch.
  • A raw user-input function parameter (named like user_input, user_message, raw_input, user_query, or message).

Safe-marker short-circuit

If a call whose name contains sanitize, clean, validate, tag_source, or mark_untrusted is applied to the value anywhere in the same function before the memory write, the finding does not fire — the short-circuit is deliberately generous, the same style as its sibling tools.

How it relates to echofence

echofence and agentmemorytaintgap are both prompt-injection-adjacent AST linters, and they are deliberately distinct, non-overlapping tools:

  • echofence — input side, single turn. Flags untrusted external content reaching a live LLM prompt directly — the indirect variant of OWASP LLM01. The risk window is one request/response cycle.
  • agentmemorytaintgap — persistence side, every future turn. Flags untrusted content being written into long-term memory that will be read back and replayed as trusted context across every subsequent turn, potentially for the lifetime of the memory store. The artifact, the timing, and the risk shape are different: a poisoned prompt affects one answer; a poisoned memory write affects all future answers until someone notices and purges the store.

See DETAILS.md for the full, honest comparison — including why this is not just "echofence but for a different sink."

Install

pip install agentmemorytaintgap

Usage

agentmemorytaintgap agent/                # scan a directory
agentmemorytaintgap memory_handler.py      # scan a file
agentmemorytaintgap agent/ --strict        # AT002 warnings fail the run too
agentmemorytaintgap agent/ --json          # machine-readable output

In CI

- run: pip install agentmemorytaintgap
- run: agentmemorytaintgap agent/ --strict

Exit codes: 0 clean · 1 a blocker (AT001), or any finding under --strict · 2 usage error.

Honest limitations

agentmemorytaintgap is a pragmatic, heuristic, single-hop, same-function-scope analyzer — not full data-flow / taint analysis. See DETAILS.md for the complete breakdown, but concretely:

  • It only recognizes the memory-write shapes and naming conventions listed above. A memory variable that doesn't contain memory/mem_store/ long_term in its name, or a tool call that doesn't match the recognized @tool/.run()/.invoke() shapes, is invisible to v0.1.
  • Tracing is single-hop and scoped to one function. A value laundered through a helper function it doesn't look inside of, or passed across functions before being written to memory, will not be traced.
  • A sanitize/tag call anywhere in the function short-circuits the finding — it trusts that the call actually does what its name implies; it does not verify that.

Treat it as a fast reviewer that catches the obvious, high-value cases on every PR, paired with human judgment for the rest.

License

MIT © Jay Tank

About

📦 Flags AI agent memory writes (LangChain, vector-store-as-memory) that persist raw, unsanitized tool/external output as trusted long-term memory - a distinct agent-memory-poisoning risk from single-turn prompt injection.

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages