fix(memory): decouple MemoryMiddleware from knowledge_store#18
Merged
Conversation
Addresses bug #3 from v0.2.0 smoke test: MemoryMiddleware was silently skipped when knowledge: false (the default config), so session memory never worked out of the box. - graph/agent.py: drop `and knowledge_store` from activation guard — memory middleware now activates whenever memory: true, regardless of knowledge store - graph/middleware/memory.py: knowledge_store is now optional (default None); guard knowledge-extraction block when store is None; add standalone prior_sessions injection via before_model when running without KnowledgeMiddleware (no double-injection: only fires when self._store is None) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
WalkthroughThe changes introduce a standalone operation mode for Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@graph/agent.py`:
- Around line 32-33: When constructing MemoryMiddleware, ensure it receives None
for the store when knowledge middleware is disabled: change the logic around
config.memory_middleware/config.knowledge_middleware so that MemoryMiddleware is
instantiated with knowledge_store only if config.knowledge_middleware is true,
otherwise pass None (e.g., MemoryMiddleware(None)). Update the middleware
creation site that currently calls MemoryMiddleware(knowledge_store) so it
conditionally passes None when config.knowledge_middleware is false, preserving
the intended memory=true, knowledge=false behavior and preventing store-backed
paths from being taken.
In `@graph/middleware/memory.py`:
- Around line 176-240: The duplicate JSON-scan→XML-format→token-budget logic in
MemoryMiddleware._load_prior_sessions and KnowledgeMiddleware.load_memory should
be extracted into a single helper function (e.g., build_prior_sessions_xml)
placed in a shared module (for example graph.middleware.shared_memory or
graph.middleware.utils). Implement the helper to accept the directory path
(MEMORY_PATH), max sessions/count limits, and token budget parameters and return
the XML string exactly matching the current outputs ("<prior_sessions/>" or
"<prior_sessions>...</prior_sessions>"). Replace
MemoryMiddleware._load_prior_sessions and KnowledgeMiddleware.load_memory to
call this helper, preserving existing behavior (file filtering, JSON decoding
fallback, message/content truncation lengths, session ordering and dropping
oldest entries to meet the token budget). Update imports and remove duplicated
code blocks from both files.
- Around line 216-229: The code builds XML fragments unsafely (ts, sid, role,
content, final) which can break the XML or alter structure; fix by XML-escaping
all inserted values (timestamp, session_id, message content, final_output) using
a standard escaper (e.g., xml.sax.saxutils.escape or html.escape) and stop using
the raw role as an element name—emit a safe tag such as <message
role="...">...</message> (escape role when used as an attribute) instead of f"
<{role}>...</{role}>", and update the code paths that create lines (the
variables ts, sid, msgs loop, role, content, final) to apply escaping
consistently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: fdf17714-1588-4bfc-966b-753a870696bf
📒 Files selected for processing (2)
graph/agent.pygraph/middleware/memory.py
Comment on lines
+32
to
33
| if config.memory_middleware: | ||
| middleware.append(MemoryMiddleware(knowledge_store)) |
There was a problem hiding this comment.
Pass None to MemoryMiddleware when knowledge middleware is off.
Line 33 still forwards a non-null knowledge_store even when config.knowledge_middleware is false. In that configuration, MemoryMiddleware takes its store-backed path, so standalone <prior_sessions> injection is skipped and post-agent extraction is turned back on. That breaks the PR’s stated memory=true, knowledge=false behavior whenever a store object is present.
Proposed fix
- if config.memory_middleware:
- middleware.append(MemoryMiddleware(knowledge_store))
+ if config.memory_middleware:
+ memory_store = knowledge_store if config.knowledge_middleware else None
+ middleware.append(MemoryMiddleware(memory_store))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@graph/agent.py` around lines 32 - 33, When constructing MemoryMiddleware,
ensure it receives None for the store when knowledge middleware is disabled:
change the logic around config.memory_middleware/config.knowledge_middleware so
that MemoryMiddleware is instantiated with knowledge_store only if
config.knowledge_middleware is true, otherwise pass None (e.g.,
MemoryMiddleware(None)). Update the middleware creation site that currently
calls MemoryMiddleware(knowledge_store) so it conditionally passes None when
config.knowledge_middleware is false, preserving the intended memory=true,
knowledge=false behavior and preventing store-backed paths from being taken.
Comment on lines
+176
to
+240
| def _load_prior_sessions(self) -> str: | ||
| """Lazy-load prior session summaries when standalone (no KnowledgeMiddleware). | ||
|
|
||
| When KnowledgeMiddleware is also in the chain it owns `<prior_sessions>` | ||
| injection. This method runs only when `self._store is None`, so there is | ||
| no double-injection risk. | ||
|
|
||
| Reads from MEMORY_PATH, returns an XML block or empty string on first | ||
| run. Mirrors KnowledgeMiddleware.load_memory() but without the store | ||
| dependency — single source of truth would be cleaner but would couple | ||
| the two files. | ||
| """ | ||
| if not os.path.isdir(MEMORY_PATH): | ||
| return "" | ||
| try: | ||
| entries = [] | ||
| for fname in os.listdir(MEMORY_PATH): | ||
| if not fname.endswith(".json"): | ||
| continue | ||
| fpath = os.path.join(MEMORY_PATH, fname) | ||
| try: | ||
| entries.append((os.path.getmtime(fpath), fpath)) | ||
| except OSError: | ||
| continue | ||
| entries.sort(reverse=True) | ||
| except OSError: | ||
| return "" | ||
| if not entries: | ||
| return "<prior_sessions/>" | ||
| summaries = [] | ||
| for _, fpath in entries[:10]: | ||
| try: | ||
| with open(fpath, encoding="utf-8") as fh: | ||
| summaries.append(json.load(fh)) | ||
| except (OSError, json.JSONDecodeError, ValueError): | ||
| continue | ||
| if not summaries: | ||
| return "<prior_sessions/>" | ||
| lines_out = [] | ||
| for s in summaries: | ||
| ts = s.get("timestamp", "unknown") | ||
| sid = s.get("session_id", "unknown") | ||
| lines = [f'<session id="{sid}" timestamp="{ts}">'] | ||
| msgs = s.get("messages", []) or [] | ||
| if msgs: | ||
| lines.append(" <messages>") | ||
| for m in msgs: | ||
| role = m.get("role", "unknown") | ||
| content = (m.get("content", "") or "")[:500] | ||
| lines.append(f" <{role}>{content}</{role}>") | ||
| lines.append(" </messages>") | ||
| final = (s.get("final_output") or "")[:300] | ||
| if final: | ||
| lines.append(f" <final_output>{final}</final_output>") | ||
| lines.append("</session>") | ||
| lines_out.append("\n".join(lines)) | ||
| # 2K token budget — chars // 4 approx, drop oldest first | ||
| while lines_out: | ||
| joined = "\n".join(lines_out) | ||
| if max(1, len(joined) // 4) <= 2000: | ||
| break | ||
| lines_out.pop() | ||
| if not lines_out: | ||
| return "<prior_sessions/>" | ||
| return "<prior_sessions>\n" + "\n".join(lines_out) + "\n</prior_sessions>" |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Extract prior-session loading into a shared helper.
This is now a second implementation of the same JSON-scan → XML-format → token-budget logic that already exists in graph/middleware/knowledge.py:61-165. Keeping both copies aligned will be error-prone, and any future change to truncation or formatting can silently make standalone memory behave differently from knowledge-backed memory.
🧰 Tools
🪛 Ruff (0.15.10)
[warning] 176-176: Too many branches (16 > 12)
(PLR0912)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@graph/middleware/memory.py` around lines 176 - 240, The duplicate
JSON-scan→XML-format→token-budget logic in MemoryMiddleware._load_prior_sessions
and KnowledgeMiddleware.load_memory should be extracted into a single helper
function (e.g., build_prior_sessions_xml) placed in a shared module (for example
graph.middleware.shared_memory or graph.middleware.utils). Implement the helper
to accept the directory path (MEMORY_PATH), max sessions/count limits, and token
budget parameters and return the XML string exactly matching the current outputs
("<prior_sessions/>" or "<prior_sessions>...</prior_sessions>"). Replace
MemoryMiddleware._load_prior_sessions and KnowledgeMiddleware.load_memory to
call this helper, preserving existing behavior (file filtering, JSON decoding
fallback, message/content truncation lengths, session ordering and dropping
oldest entries to meet the token budget). Update imports and remove duplicated
code blocks from both files.
Comment on lines
+216
to
+229
| ts = s.get("timestamp", "unknown") | ||
| sid = s.get("session_id", "unknown") | ||
| lines = [f'<session id="{sid}" timestamp="{ts}">'] | ||
| msgs = s.get("messages", []) or [] | ||
| if msgs: | ||
| lines.append(" <messages>") | ||
| for m in msgs: | ||
| role = m.get("role", "unknown") | ||
| content = (m.get("content", "") or "")[:500] | ||
| lines.append(f" <{role}>{content}</{role}>") | ||
| lines.append(" </messages>") | ||
| final = (s.get("final_output") or "")[:300] | ||
| if final: | ||
| lines.append(f" <final_output>{final}</final_output>") |
There was a problem hiding this comment.
Escape persisted values before embedding them in <prior_sessions>.
These fields are inserted raw into XML. A saved message containing <, &, or tag-like text will produce malformed prompt context on the next run, and using role as an element name also lets malformed session files change the structure of the injected block.
Proposed fix
+ from xml.sax.saxutils import escape
lines_out = []
for s in summaries:
- ts = s.get("timestamp", "unknown")
- sid = s.get("session_id", "unknown")
+ ts = escape(str(s.get("timestamp", "unknown")))
+ sid = escape(str(s.get("session_id", "unknown")))
lines = [f'<session id="{sid}" timestamp="{ts}">']
msgs = s.get("messages", []) or []
if msgs:
lines.append(" <messages>")
for m in msgs:
- role = m.get("role", "unknown")
- content = (m.get("content", "") or "")[:500]
- lines.append(f" <{role}>{content}</{role}>")
+ role = escape(str(m.get("role", "unknown")))
+ content = escape(str((m.get("content", "") or "")[:500]))
+ lines.append(f' <message role="{role}">{content}</message>')
lines.append(" </messages>")
- final = (s.get("final_output") or "")[:300]
+ final = escape(str((s.get("final_output") or "")[:300]))
if final:
lines.append(f" <final_output>{final}</final_output>")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@graph/middleware/memory.py` around lines 216 - 229, The code builds XML
fragments unsafely (ts, sid, role, content, final) which can break the XML or
alter structure; fix by XML-escaping all inserted values (timestamp, session_id,
message content, final_output) using a standard escaper (e.g.,
xml.sax.saxutils.escape or html.escape) and stop using the raw role as an
element name—emit a safe tag such as <message role="...">...</message> (escape
role when used as an attribute) instead of f" <{role}>...</{role}>", and
update the code paths that create lines (the variables ts, sid, msgs loop, role,
content, final) to apply escaping consistently.
This was referenced Apr 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes bug #3 from v0.2.0 smoke testing: MemoryMiddleware silently skipped when knowledge: false (the default config).
Root cause
graph/agent.py:32 had an
and knowledge_storeconjunction on the MemoryMiddleware activation guard. Default config hasknowledge: falsesoknowledge_store=None, silently disabling memory despitememory: true.Changes
and knowledge_storefrom the guard<prior_sessions>injection viabefore_model— only fires whenself._store is None, so no double-injection with KnowledgeMiddleware.Behavior by config
Authored by Ava after the original agent completed the plan correctly but hit PR-create rate limits 3x and the retry budget exhausted.
Summary by CodeRabbit