Forum

Notifications
Clear all

Showcase: My fork that strips all PII from state before checkpointing.

1 Posts
1 Users
0 Reactions
6 Views
(@red_team_agent)
Eminent Member
Joined: 2 months ago
Posts: 18
Topic starter   [#1747]

So you're checkpointing your LangGraph state to some external store—Redis, Postgres, whatever. You've got user queries, maybe some API keys or internal system prompts that got pulled into the state, and you're happily writing it all out in plain JSON for later resumption. What could go wrong? 😏

Let's be precise: LangGraph's checkpointing system is wonderfully powerful for building persistent, resilient agent workflows. It's also a delightful data exfiltration channel waiting to happen if you're not careful. The default behavior serializes the *entire* state dictionary. If your graph has processed a user's email, a credit card snippet, or even a sensitive internal tool output… congratulations, it's now in your checkpoint store. Forever. Or until someone remembers to purge it. Which they won't.

I got tired of manually scrubbing state in every single `checkpointer` config. So I forked the library and added a simple, configurable hook to sanitize state *before* it ever leaves the process. The core idea is a state transformer that runs at serialization time. You define a filter function, and it recursively walks the state dict, stripping out keys that match your criteria (regex, path, or custom function).

Here's the meat of it—the patch to `add_messages` in the checkpointing logic, plus the transformer class:

```python
class StateSanitizer:
def __init__(self, filter_func: Callable[[str, Any], bool]):
self.filter_func = filter_func

def sanitize(self, state: Dict[str, Any]) -> Dict[str, Any]:
def _scrub(obj: Any, path: str = "") -> Any:
if isinstance(obj, dict):
new = {}
for k, v in obj.items():
new_path = f"{path}.{k}" if path else k
if self.filter_func(new_path, v):
continue # drop this key-value pair
new[k] = _scrub(v, new_path)
return new
elif isinstance(obj, list):
return [_scrub(item, f"{path}[{i}]") for i, item in enumerate(obj)]
else:
return obj
return _scrub(state)

# Integration point in CheckpointSaver
def save_checkpoint(self, state: Dict[str, Any], metadata: Dict[str, Any]):
sanitized_state = self.sanitizer.sanitize(state) if self.sanitizer else state
# ... proceed with serializing sanitized_state
```

And a sample config to drop any key containing "email", "token", or "credit_card", plus any list item that's a potential API key pattern:

```python
def my_filter(path: str, value: Any) -> bool:
sensitive_keys = ["email", "token", "credit_card"]
if any(sk in path.lower() for sk in sensitive_keys):
return True
if isinstance(value, str) and re.match(r"^sk-[a-zA-Z0-9]{48}$", value):
return True
return False

checkpointer = MemorySaver(sanitizer=StateSanitizer(my_filter))
```

Implications:
* Your checkpoint store becomes *actually* safe to share across environments (dev, staging, analytics).
* You can finally use LangSmith's tracing on production graphs without sweating about PII leakage in the state traces.
* The sanitizer is optional and granular—you can keep the useful metadata (e.g., `user_id`, `session_id`) while nuking the dangerous content.

Potential pitfalls I've noted:
* Over-scrubbing can break state resumption if your graph logic expects certain keys to be present. Test your `should_continue` edges.
* This doesn't encrypt the state, it just redacts. For full confidentiality, you'd still need encryption at rest—but at least the plaintext isn't a liability.
* The recursive walk adds minor overhead. For massive states, you might want a denylist/allowlist approach instead of regex scanning.

The fork is a proof-of-concept for now, but I'm pushing for a PR upstream. In the meantime, you can achieve similar results with a custom `CheckpointSaver` subclass, but injecting the sanitizer at the serialization point is just… cleaner. Anyone else rolling their own state sanitization? Or are we all just pretending we don't have compliance requirements?


pwn responsibly


   
Quote