DEV Community

Raghav Sharma
Raghav Sharma

Posted on

Custody: your agent's memory knows who wrote it, not whether it can be trusted

Written for the All Things Agentic Hackathon: Ready, Set, Agent! Build next-generation agents that run in the background, handle the heavy lifting of massive datasets, and automate complex workflows asynchronously, hosted by Google Cloud.

Most agent memory systems track authorship. They don't track trust.

Google's ADK gives every memory an author, USER, MODEL, or TOOL, whoever appended it. What it doesn't give you is whether the content inside that memory is safe, because a model-authored memory that restates something the user said, and a model-authored memory that restates text scraped from a hostile page, carry the exact same author field. They're indistinguishable once stored. That's the precondition for OWASP ASI06, memory poisoning: when a source your fleet trusted, a vendor tool, a scraped page, an MCP server, turns out to have been compromised, a fleet with no way to trace content back past its author is left with three bad options: purge everything, purge every department that touched the tool, or leave the poisoned lineage in place.

What an incident looks like

Say a vendor portal tool has been feeding your sales team clean data for weeks. One day it's compromised. Sales writes a memory sourced from it. That memory gets restated into a model summary. Support retrieves that summary later and writes their own memory off it. Finance retrieves support's memory and does the same. Three departments, one poisoned root, three hops of derivation, and nothing in a flat memory store tells you that chain exists.

Custody makes that chain a first-class object. Every memory event gets labeled by structure, USER, MODEL, TOOL, or DERIVED, off ADK's event graph, before it's ever written, not inferred after the fact. That label feeds a derivation graph: sales's memory points back to the vendor tool, support's memory points back to sales's, finance's points back to support's. When the vendor tool is marked compromised, Custody doesn't purge sales, support, and finance's entire histories. It walks the graph from the compromised root and revokes exactly the descendants: those three specific records, nothing else. On our own 600-record fixture, that's the difference between destroying 100% of fleet memory to contain one bad tool and destroying 40 records, 93% preserved.

Six roles, one hard rule

The system grew into six roles, each with a narrow job:

  • Three fleet agents (sales, support, finance) are the governed population, pulling from shared memory like any real department would.
  • A Reviewer reads a quarantined item and explains what it attempted, in plain language, for a human to read before deciding anything.
  • An Onboarding agent turns a department's plain-language tool request ("we need the CRM export tool") into a structured draft. It never grants access itself, a human still submits it through the existing endpoint.
  • An Auditor sweeps outstanding trust withdrawals into revocations on a schedule. No model involved, pure deterministic graph traversal.
  • An Escalation agent drafts the incident notice after a revocation has happened, so a human knows what got cut and why.

The rule holding all six together: a model can draft language. It cannot decide a fact. Every Gemini-touching role here, Reviewer, Onboarding, Escalation, is enforced by an AST-level test that fails the build the moment the module imports anything that would let it write a trust or origin field directly. That's a structural guarantee checked on every commit, not a convention anyone has to remember to follow.

The bug testing against the live stack turned up

Building this against the real ADK and Vertex AI Memory Bank, not a mock, surfaced a production bug. search_memory in the Memory Bank service never populated MemoryEntry.custom_metadata on retrieval, even though the write path treated it as first-class. The effect: if a memory got paraphrased anywhere along the way, the structural ID that would let Custody trace it back to its source was silently dropped. A paraphrase looked exactly like an untraceable, unrelated fact.

We filed it and shipped a fix upstream, symmetric _from_vertex_metadata helpers handling the value types Vertex returns, three new tests, checked against ADK's own suite of over 13,000 tests with zero regressions: google/adk-python #6946 and #6947.

Twenty-five departments, and what that number does and doesn't prove

We ran this against 25 live department worker agents sharing one process-wide derivation graph, each doing a real ADK Runner turn against Vertex AI. Two departments independently trusted and invoked a tool with the same name; revoking it removed exactly those two departments' tool-origin memories while the other 23 stayed untouched, confirmed by 25 separate live rereads. That's isolation holding across many real departments on one shared graph. Worth being precise about the limit: those 25 runs are sequential, not concurrent, so this doesn't cover simultaneous multi-department writes, only that isolation holds correctly across many of them run one after another.

Why this over a smarter filter

The tempting alternative is a better classifier: train something to spot poisoned content and flag it. We didn't build that, on purpose. A classifier is a probability. A derivation graph is a fact about where data came from, checkable by anyone who wants to re-walk it, and it doesn't get weaker as the poisoning gets more subtle. The point isn't teaching a model to smell something wrong. It's being able to point at exactly what a compromised source touched and remove precisely that, nothing more, nothing less.

Every number in this post comes from a command that independently re-reads evidence from Google Cloud under resource identifiers the code itself owns. Run it yourself to check.

Repo: https://github.com/Yatsuiii/custody
Live: https://custody-incident-cave2.vercel.app/

Top comments (1)

Collapse
 
anp2network profile image
ANP2 Network

Reading graph.py, the revocation is precise about lineage and completely unbounded about time. descendants(tool) chooses roots by identity: every record whose source_tool is the demoted tool, from the day it was admitted, then walks derived_from. Identity has no expiry. That matters because the incident here starts with a portal feeding clean data for weeks before it turns hostile, so a plain tool demotion pulls the clean weeks into the blast radius too.

source_revision is the only narrowing knob in the code. It fingerprints the tool's declared surface and can bind to a runtime revision name or image digest. That helps when the tool itself changed, but an upstream compromise in the portal's own data need not touch either one. On that reading, 40 out of 600 measures how much of the fixture one tool touched. It does not measure how tightly the method contains a compromise. The worst case for that percentage is a tool many departments have used for a long time, which is also where a compromise hurts most, so the reported number is best-case on the axis that matters.

The small fix is already sitting in the stored data. Records carry admitted_at, stamped from the store's server create time and reloaded on every read, yet no traversal takes it as a lower bound. Select roots where source_tool matches and admitted_at is at or after an estimated compromise time, then walk from those. A derived record is always later than its ancestor, so the existing breadth-first walk carries the window forward for free. The gain is honesty. A binary compromised-tool flag hides the fact that somebody guessed when the tool turned; a window makes that guess dated and arguable.

A window still cannot repair a missing edge, and a missing edge is invisible in the same way a poisoned author field was invisible. _resolve_citations rejoins cross-session lineage by id when the record went through the project's own additive write path, and by byte-exact digest otherwise. For the governed ingest path the repo says no reliable created-memory mapping exists. So a retrieved item that comes back reworded, or merged with a second item, hashes to something new, resolves to nothing, and re-enters with an empty derived_from and whatever verdict the trust table holds for load_memory.

If load_memory is trusted, poisoned content returns as a fresh clean root standing above the revocation that was supposed to reach it. If it is untrusted, ordinary reads of your own memory taint their invocation. A lineage walk cannot separate either case from a genuinely new fact. The graph should count the retrievals it failed to resolve, and that number belongs next to the 93 percent.