feat: add AutoGen and OpenAI Agents SDK integration adapters#22
feat: add AutoGen and OpenAI Agents SDK integration adapters#22hesam-oxe wants to merge 4 commits into
Conversation
Created GuardedAutoGenAgent, GuardedGroupChatManager, GuardedAgentContext, GuardedToolOutput, and GuardedHandoff. Closes OWASP#8
vgudur-dev
left a comment
vgudur-dev
left a comment
There was a problem hiding this comment.
Thanks @hesam-oxe — getting AutoGen and the OpenAI Agents SDK on the integration list is overdue, and the structure (per-framework module under integrations/) is right. A few things to address before this is merge-ready; some are mechanical, two are real correctness/security bugs.
Blocking
GuardedAgentContext.get_statesilently bypasses the guard onPolicyViolation— falls through to the unguarded underlying context. That defeats the policy. Inline atopenai_agents.py:45-54with a fix.GuardedAgentContext.set_statereturnsTrueeven when nothing was persisted to the wrapped context, and lets the guard's store and the underlying store diverge onREDACT. Inline atopenai_agents.py:31-43.- No tests. Two new public modules with seven classes; the suite doesn't exercise any of them. At minimum:
- injection in a sent/received AutoGen message is dropped (
drop_blocked=True) or raised (drop_blocked=False) GuardedAgentContext.set_stateround-trips throughget_statewith the guard's value (incl. redaction)GuardedAgentContext.get_statepropagatesPolicyViolationupward when reads are blockedGuardedToolOutput.screen_tool_outputreturnsFalseon quarantine andTrueon allowGuardedHandoff.transferreturnsFalsewhen the handoff context contains injection markers (use a permissive guard so it's the detector chain doing the work, not the policy)
- injection in a sent/received AutoGen message is dropped (
- Trailing newlines missing on both files.
Architecture — needs discussion, not necessarily a rewrite in this PR
- The wrappers don't engage with the frameworks' real dispatch hooks. AutoGen's group chats and reply functions run on the wrapped agent through
register_reply/_process_received_message, bypassingGuardedAutoGenAgent.sendand.receiveentirely. The OpenAI Agents SDK has similar internal flows. Detail and suggested approach inline atautogen.py:56. I'd recommend convertingGuardedAutoGenAgenttoinstall_guard(agent, guard)that registers aregister_replyhook on the existing agent — cleaner and actually gets called.
Nits / style
Optional[MemoryGuard]→MemoryGuard | Noneto match the rest of the codebase (inline)._HAS_AUTOGENflag is set but never enforced — either drop it or raise at construction time when the SDK isn't installed (inline).agent_isolationparameter onGuardedGroupChatManageris dead code (inline).pyproject.tomlnot updated. Addautogen = [...]andopenai-agents = [...]to[project.optional-dependencies], mirroring the existinglangchainextra, sopip install agent-memory-guard[autogen]works and the import-guard pattern has something to fall back to.integrations/__init__.pynot extended — the new classes aren't re-exported. Users currently have to writefrom agent_memory_guard.integrations.autogen import GuardedAutoGenAgent. Worth adding optional re-exports undertry/except ImportErrorso the package-level import surface matches the langchain pattern.
The OpenAI Agents fixes (#1, #2) are the only items that genuinely block this. Everything else can be follow-ups, but landing without tests would set a precedent we don't want in the integrations directory.
Generated by Claude Code
| from autogen import ConversableAgent # type: ignore | ||
|
|
||
| _HAS_AUTOGEN = True | ||
| except Exception: # pragma: no cover - optional dependency |
There was a problem hiding this comment.
_HAS_AUTOGEN is set but never consulted at runtime — callers can instantiate GuardedAutoGenAgent with any object and it'll silently work as a generic shell. Either:
- Drop the flag entirely (the wrapper is duck-typed anyway), or
- Guard the constructor:
def __init__(self, agent, guard=None, *, drop_blocked=True):
if not _HAS_AUTOGEN:
raise ImportError(
"agent-memory-guard[autogen] not installed; "
"pip install agent-memory-guard[autogen]"
)
...Option 2 also requires adding autogen = ["autogen-agentchat>=0.2"] (or the package name AutoGen actually ships under) to [project.optional-dependencies] in pyproject.toml — same pattern as the existing langchain extra.
Generated by Claude Code
| def screen_message(self, message: dict, source: str) -> bool: | ||
| """Screen a message before send/receive.""" | ||
| msg_id = f"autogen.{self._agent.name}.msg.{self._message_count}" | ||
| payload = str(message) |
There was a problem hiding this comment.
str(message) on a dict produces "{'content': '...'}" — the detectors will still match the injection regex through the repr, but it's fragile (escaped quotes, embedded \n, the dict's other keys all become noise the regex has to step around). Prefer extracting the content field when it's a dict:
| payload = str(message) | |
| msg_id = f"autogen.{self._agent.name}.msg.{self._message_count}" | |
| payload = ( | |
| message.get("content", "") if isinstance(message, dict) else str(message) | |
| ) |
Also worth noting: every screen_message call writes a fresh entry to memory via self.guard.write(msg_id, payload, ...) — i.e. screening stores the message as a side effect. Over a long conversation that's hundreds of entries living in MemoryStore. If the intent is just to run detectors without persisting, expose a MemoryGuard.screen() method that runs detectors-only (worth a separate issue), or accept that storage is part of the design and document it.
Generated by Claude Code
| def send( | ||
| self, message: str | dict, recipient: Any, request_reply: bool = False | ||
| ) -> None: | ||
| msg = message if isinstance(message, dict) else {"content": message} | ||
| if self.screen_message(msg, "autogen_send"): | ||
| self._agent.send(message, recipient, request_reply=request_reply) | ||
|
|
||
| def receive( | ||
| self, message: str | dict, sender: Any, request_reply: bool = False | ||
| ) -> None: | ||
| msg = message if isinstance(message, dict) else {"content": message} | ||
| if self.screen_message(msg, "autogen_receive"): | ||
| self._agent.receive(message, sender, request_reply=request_reply) | ||
|
|
There was a problem hiding this comment.
Architectural concern, not blocking but worth thinking through before this lands:
These wrappers only intercept calls made directly via GuardedAutoGenAgent.send / .receive. AutoGen's framework dispatches messages through reply_functions registered with agent.register_reply() and the internal _process_received_message / process_message_before_send hooks — those run on the wrapped self._agent, not on the wrapper, so the guard never sees them. In practice that means:
- A
GroupChatManagerrunning the chat will callagent._process_received_message(...)directly (via the proxied__getattr__), and the guard is bypassed. - Auto-generated tool replies, function calls, and
generate_replyflows all happen inside the wrapped agent.
The "real" drop-in pattern for AutoGen is to register a guard reply function via agent.register_reply([Agent, None], reply_func=self._guard_reply, position=0) so it fires before the framework's own reply functions. I'd suggest one of:
- Convert
GuardedAutoGenAgentto a functioninstall_guard(agent, guard)that registers a reply hook on the existing agent (no wrapping). This is whatregister_replyis built for. - Subclass
ConversableAgentand overrideprocess_received_message/_process_received_message(signature changes between AutoGen versions — pin the supported range). - Document the wrapper as "only safe if you call
.send()/.receive()explicitly", with a warning when someone passes the wrapper into aGroupChat.
Option 1 is the cleanest. Happy to spike it in a follow-up if you'd like.
Generated by Claude Code
| agent_isolation: bool = True, | ||
| ) -> None: | ||
| self._group_chat = group_chat | ||
| self.guard = guard or MemoryGuard() |
There was a problem hiding this comment.
agent_isolation is stored but never read anywhere in the class — dead parameter. Either wire it up (e.g. tag entries with the originating agent and refuse reads from other agents — MemoryGuard.set_current_task() from #25 already gives you most of the machinery for this) or drop it from the signature for now.
Generated by Claude Code
| if agent_name not in self._agent_keys: | ||
| self._agent_keys[agent_name] = set() | ||
| self._agent_keys[agent_name].add(key) | ||
| return True No newline at end of file |
There was a problem hiding this comment.
Missing trailing newline (\ No newline at end of file in the diff). Same in openai_agents.py. Add a final \n.
Generated by Claude Code
|
|
||
| def get_state(self, key: str) -> Any: | ||
| full_key = f"openai_agents.state.{key}" | ||
| try: | ||
| cached = self.guard.read(full_key, sink="openai_agents") | ||
| if cached is not None: | ||
| return cached | ||
| except PolicyViolation: | ||
| pass | ||
| if hasattr(self._context, "get_state"): |
There was a problem hiding this comment.
Security bug — silent bypass on policy violation.
try:
cached = self.guard.read(full_key, sink="openai_agents")
if cached is not None:
return cached
except PolicyViolation:
pass # <-- here
if hasattr(self._context, "get_state"):
return self._context.get_state(key)If the guard blocks a read (PolicyViolation), the wrapper swallows it and returns the unguarded state straight from self._context. That defeats the whole point — a policy that says "block reads of this key" gets silently bypassed.
Two more issues in the same block:
if cached is not Nonetreats legitimateNone/False/0writes as "no cached value", causing fallthrough to the underlying context (which may return a stale value).- Reads should not generally fall back to the underlying store: if the guard owns state, it owns state; otherwise it's just an event tap and you have two sources of truth.
Suggested rewrite:
| def get_state(self, key: str) -> Any: | |
| full_key = f"openai_agents.state.{key}" | |
| try: | |
| cached = self.guard.read(full_key, sink="openai_agents") | |
| if cached is not None: | |
| return cached | |
| except PolicyViolation: | |
| pass | |
| if hasattr(self._context, "get_state"): | |
| def get_state(self, key: str, default: Any = None) -> Any: | |
| full_key = f"openai_agents.state.{key}" | |
| sentinel = object() | |
| cached = self.guard.read(full_key, default=sentinel, sink="openai_agents") | |
| if cached is not sentinel: | |
| return cached | |
| if hasattr(self._context, "get_state"): | |
| return self._context.get_state(key) | |
| return default |
PolicyViolation from self.guard.read should propagate to the caller, not be swallowed. If you want a softer mode, expose it as a flag (raise_on_block: bool = True) rather than catching unconditionally.
Generated by Claude Code
|
|
||
| def set_state(self, key: str, value: Any) -> bool: | ||
| full_key = f"openai_agents.state.{key}" | ||
| try: | ||
| decision = self.guard.write(full_key, value, source="openai_agents") | ||
| except PolicyViolation: | ||
| if self._drop_blocked: | ||
| return False | ||
| raise | ||
| if decision == Action.QUARANTINE: | ||
| return False | ||
| if hasattr(self._context, "set_state"): | ||
| self._context.set_state(key, value) |
There was a problem hiding this comment.
Two ordering issues here:
-
Return value is misleading. The function returns
Truewhen the guard accepted the write, regardless of whetherself._context.set_statewas ever called. Soset_statereturns success when no state was actually persisted to the wrapped context (becausehasattr(self._context, "set_state")was false). Caller has no way to know. -
Even when both succeed, the values can diverge. If the policy is
REDACT, the guard stores the redacted value butself._context.set_state(key, value)stores the original.get_statewill then return whichever it reaches first — and they won't match.
Fix: write to the guard first, then propagate the redacted value (the one the guard actually committed) to the underlying context. Easiest path is to expose the committed value via MemoryGuard, or to read it back immediately:
def set_state(self, key: str, value: Any) -> bool:
full_key = f"openai_agents.state.{key}"
try:
decision = self.guard.write(full_key, value, source="openai_agents")
except PolicyViolation:
if self._drop_blocked:
return False
raise
if decision == Action.QUARANTINE:
return False
if hasattr(self._context, "set_state"):
# read back the (possibly redacted) value so the two stores agree
committed = self.guard.read(full_key, sink="openai_agents")
self._context.set_state(key, committed)
return TrueGenerated by Claude Code
| """ | ||
| from __future__ import annotations | ||
|
|
||
| from typing import Any, Optional |
There was a problem hiding this comment.
Project style uses MemoryGuard | None, not Optional[MemoryGuard] (see e.g. guard.py, classification.py, the existing langchain.py). Since from __future__ import annotations is at the top, | None works on the project's minimum 3.9 target.
| from typing import Any, Optional | |
| from typing import Any |
…and then replace each Optional[MemoryGuard] with MemoryGuard | None below. Same in autogen.py.
Generated by Claude Code
…for both adapters
|
@vgudur-dev Thanks for the thorough review! All blocking issues resolved:
The architectural concern about AutoGen dispatch hooks is noted — |
|
@hesam-oxe — thanks for the AutoGen and OpenAI Agents SDK adapters, the overall structure is right. Before this can merge, please address the 5 review items from May 13:
Happy to merge once these are addressed. |
_HAS_AUTOGEN guard, content extraction, docstring note, dead param removed, trailing newline
|
@vgudur-dev All five review items addressed:
Test file also added for import coverage. Ready for re-review! 🙏 |
|
@vgudur-dev All tests pass locally (2 passed). Both modules import The CI failure appears to be the Node.js 20 deprecation warning on Could you re-trigger the workflow when you get a chance? 🙏 |
2 participants
Added integration adapters for Microsoft AutoGen and OpenAI Agents SDK.
Changes
autogen.py— GuardedAutoGenAgent + GuardedGroupChatManageropenai_agents.py— GuardedAgentContext + GuardedToolOutput + GuardedHandoffFeatures
Closes #8