feat: add per-request isolation headers for multi-tenant support #29
Conversation
Implements X-Graph-Name and X-Collection-Name headers to enable per-request graph and collection routing without deploying separate AutoMem instances. Key changes: - Add _get_graph_name() and _get_collection_name() helper functions with validation (regex, length limits) and optional whitelist support - Update get_memory_graph() to support per-request isolation with automatic header detection and connection caching - Update core endpoints (store_memory, update_memory, delete_memory, recall_memories) to use per-request isolation - Add comprehensive test suite (22 tests) covering header extraction, validation, whitelists, and backwards compatibility - Maintain full backwards compatibility: no headers = environment defaults - Add connection caching to avoid redundant graph instance creation Security features: - Header validation: [a-zA-Z0-9_-]+ with max 64 chars - Optional whitelists via ALLOWED_GRAPHS and ALLOWED_COLLECTIONS - 400 for invalid format, 403 for whitelist rejection Implementation follows spec in: services/automem-federation/docs/automem-isolation-headers-spec.md Closes: Multi-tenant isolation requirement
- Import has_request_context from Flask - Add context check to _get_graph_name() and _get_collection_name() - Background tasks (consolidation, enrichment) now safely use defaults
📝 WalkthroughWalkthroughThis PR introduces per-request graph and collection isolation through HTTP headers ( Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
app.py (1)
1893-1937: Async embedding pipeline ignores per-request graph/collection for header-based writes
_store_embedding_in_qdrantcurrently:
- Calls
graph = get_memory_graph()with nograph_name, which resolves toGRAPH_NAMEoutside a request context.- Always writes to Qdrant using
collection_name=COLLECTION_NAME.In
store_memory, for requests withX-Graph-Name/X-Collection-Nameset, you:
- Write the memory into
graph = get_memory_graph(graph_name).- Queue async embedding via
enqueue_embedding(memory_id, content)when no embedding is provided.When the background worker later calls
_store_embedding_in_qdrant, it will:
- Look for the memory in the default graph, not the header-selected graph (likely not found).
- Write (or skip writing) vectors only in the default collection, regardless of the original per-request collection.
Net effect: header-isolated memories saved without an explicit embedding will report
embedding_status="queued"but may never get embedded/stored in the intended tenant collection.Consider threading isolation context through the embedding pipeline, for example:
- Extend
enqueue_embedding(and the queued payload) to includegraph_nameand/orcollection_name.- Pass those through to
_store_embedding_in_qdrant, and callget_memory_graph(graph_name)plus_get_collection_name()(or a passed-in value) there.This keeps background workers defaulting to env-based names when called without explicit context, while allowing header-based writes to behave consistently.
🧹 Nitpick comments (9)
app.py (8)
1070-1073: _graph_cache is used, _collection_cache is not
_graph_cacheis correctly used byget_memory_graph, but_collection_cacheis currently unused anywhere in this module.To avoid confusion and lint noise, either remove
_collection_cachefor now or wire it up where per-collection caching is actually needed (e.g., if you later cache Qdrant connections/metadata).
1075-1099: _get_graph_name helper matches isolation/whitelist requirementsThe helper correctly:
- Falls back to
GRAPH_NAMEwhen there is no request context (background tasks).- Enforces the
[A-Za-z0-9_-]+pattern and a 64-char limit.- Applies a simple ALLOWED_GRAPHS whitelist with whitespace-tolerant parsing.
This behavior looks sound for per-request graph isolation. If you ever hit performance issues from repeated env parsing, you could cache the parsed whitelist once at module import, but that’s purely an optimization.
1101-1125: _get_collection_name mirrors graph behavior correctly
_get_collection_namemirrors_get_graph_name:
- Safe outside request context (returns
COLLECTION_NAME).- Enforces the same regex/length constraints.
- Implements ALLOWED_COLLECTIONS with whitespace-trimming.
This is consistent and should be enough for per-request Qdrant collection isolation. Same as above, caching the parsed whitelist would only be a micro-optimization.
1439-1471: get_memory_graph caching and header handling are mostly good, with minor cleanupsThe new
get_memory_graph(graph_name: Optional[str] = None)behavior looks reasonable:
- Initializes FalkorDB lazily via
init_falkordb().- Preserves the old “no-connection but pre-set
state.memory_graph” behavior for tests/edge cases whengraph_name is None.- Uses
_get_graph_name()so background tasks seeGRAPH_NAMEand request handlers pick upX-Graph-Name.- Caches
select_graphresults per-name in_graph_cache.Two small improvements:
The
try/except RuntimeErroraround_get_graph_name()can be removed:_get_graph_name()already guards onhas_request_context()and usesabort()for invalid headers, so it won’t raiseRuntimeError.If you expect many concurrent requests and graph names, consider whether
_graph_cacheneeds a simple lock orfunctools.lru_cache-style wrapper to avoid races on first-insertion. In practice CPython’s GIL usually makes this benign, but an explicit guard can make the intent clearer.
503-560: Vector tag-only search still hardcodes COLLECTION_NAME
_vector_filter_only_tag_searchalways uses the globalCOLLECTION_NAMEforscroll(). If the intent is that recall/tag-only searches should respectX-Collection-Name, you’ll need to resolve the effective collection per-request (e.g., via_get_collection_name()).A minimal adjustment would be:
def _vector_filter_only_tag_search( qdrant_client: Optional[QdrantClient], @@ - query_filter = _build_qdrant_tag_filter(tag_filters, tag_mode, tag_match) + query_filter = _build_qdrant_tag_filter(tag_filters, tag_mode, tag_match) @@ - points, _ = qdrant_client.scroll( - collection_name=COLLECTION_NAME, + collection_name = _get_collection_name() + points, _ = qdrant_client.scroll( + collection_name=collection_name, scroll_filter=query_filter, limit=limit, with_payload=True, )This will still fall back to
COLLECTION_NAMEoutside of request context due to_get_collection_name().
562-633: Vector search also hardcodes COLLECTION_NAME; consider using _get_collection_name()Similarly,
_vector_searchalways queries againstCOLLECTION_NAME, so vector recalls won’t honorX-Collection-Nameunless other layers compensate.To align with the new isolation behavior:
def _vector_search( @@ - try: - vector_results = qdrant_client.search( - collection_name=COLLECTION_NAME, + try: + collection_name = _get_collection_name() + vector_results = qdrant_client.search( + collection_name=collection_name, query_vector=embedding, limit=limit, with_payload=True, query_filter=query_filter, )This preserves current behavior for background callers (no request context) but enables per-request collection routing for recall.
2771-2776: delete_memory uses per-request isolation but may leave legacy vectors
delete_memorynow:
- Selects the graph with
graph_name = _get_graph_name()andget_memory_graph(graph_name).- Deletes the Qdrant point from
collection_name = _get_collection_name().This is correct for new header-aware traffic. Be aware that memories created before isolation support (always in the default collection) but later deleted under a different
X-Collection-Namewon’t have their legacy vectors removed. If that migration scenario matters, you may want a one-off cleanup or a fallback delete in the default collection when the tenant-specific delete finds nothing.Also applies to: 2788-2793
2884-2890: recall_memories: unused locals and redundant header extractionIn
recall_memories:graph_name = _get_graph_name() collection_name = _get_collection_name() seen_ids: set[str] = set() graph = get_memory_graph(graph_name) qdrant_client = get_qdrant_client()These locals are never used;
handle_recallis only passed the callables (get_memory_graph,get_qdrant_client, etc.) and re-resolves resources itself. Ruff correctly flagscollection_name,seen_ids, andgraphas unused.You can safely drop these assignments:
- # Extract per-request isolation headers - graph_name = _get_graph_name() - collection_name = _get_collection_name() - - seen_ids: set[str] = set() - graph = get_memory_graph(graph_name) - qdrant_client = get_qdrant_client() - - results: List[Dict[str, Any]] = [] - vector_matches: List[Dict[str, Any]] = [] + results: List[Dict[str, Any]] = [] + vector_matches: List[Dict[str, Any]] = []Since
get_memory_graphalready consults_get_graph_name()internally, isolation is still honored where used.tests/test_isolation_headers.py (1)
16-33: Unused fixtures mock_graph and mock_qdrant
mock_graphandmock_qdrantfixtures are defined but never used in this test module.If you don’t plan to reuse them from other tests, consider removing them to keep the test file focused; if they’re meant for future use, it can help to add a brief comment or move them to
conftest.pywhere shared fixtures usually live.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
app.py(11 hunks)tests/test_isolation_headers.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
**/*.py: Run 'black .' to format Python code before committing
Run 'flake8' to lint Python code for quality issues
**/*.py: Use type hints in Python code
Use 4-space indentation in Python code
Enforce maximum line length of 100 characters (Black configuration)
Format Python code with Black
Sort imports with Isort using profile=black
Lint Python code with Flake8
Use snake_case for module and function names in Python
Use PascalCase for class names in Python
Use UPPER_SNAKE_CASE for constant names in Python
Files:
app.pytests/test_isolation_headers.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
tests/**/*.py: Set environment variablePYTEST_DISABLE_PLUGIN_AUTOLOAD=1when running pytest to disable conflicting plugins
Use pytest with DummyGraph fixture to mock FalkorDB operations in unit tests
Files:
tests/test_isolation_headers.py
tests/test_*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/test_*.py: Name test files with test_ prefix and place in tests/ directory using Pytest framework
Use Pytest fixtures instead of global variables in tests
Files:
tests/test_isolation_headers.py
🧠 Learnings (3)
📚 Learning: 2025-12-09T02:15:22.291Z
Learnt from: CR
Repo: verygoodplugins/automem PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-09T02:15:22.291Z
Learning: Applies to automem/models/**/*.py : Use UUID for memory IDs and store in both graph (FalkorDB) and vector (Qdrant) databases for cross-referencing
Applied to files:
app.py
📚 Learning: 2025-12-09T02:15:22.291Z
Learnt from: CR
Repo: verygoodplugins/automem PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-09T02:15:22.291Z
Learning: Applies to tests/**/*.py : Use pytest with DummyGraph fixture to mock FalkorDB operations in unit tests
Applied to files:
app.pytests/test_isolation_headers.py
📚 Learning: 2025-12-09T02:15:22.291Z
Learnt from: CR
Repo: verygoodplugins/automem PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-09T02:15:22.291Z
Learning: Use FalkorDB connection configuration with FALKORDB_HOST, FALKORDB_PORT, and FALKORDB_GRAPH environment variables for graph database access
Applied to files:
app.py
🧬 Code graph analysis (2)
app.py (5)
tests/conftest.py (2)
select_graph(18-22)delete(49-50)automem/api/memory.py (1)
delete(412-439)consolidation.py (1)
delete(41-44)tests/test_consolidation_engine.py (1)
delete(76-77)tests/test_api_endpoints.py (1)
delete(206-212)
tests/test_isolation_headers.py (2)
app.py (3)
_get_graph_name(1075-1098)_get_collection_name(1101-1124)get_memory_graph(1439-1471)tests/conftest.py (1)
select_graph(18-22)
🪛 GitHub Actions: CI
app.py
[error] 1-1: Black formatting check reformatted code in this file.
tests/test_isolation_headers.py
[error] 1-1: Black formatting check reformatted code and isort modified imports in this file.
🪛 Ruff (0.14.8)
app.py
2886-2886: Local variable collection_name is assigned to but never used
Remove assignment to unused variable collection_name
(F841)
2888-2888: Local variable seen_ids is assigned to but never used
Remove assignment to unused variable seen_ids
(F841)
2889-2889: Local variable graph is assigned to but never used
Remove assignment to unused variable graph
(F841)
tests/test_isolation_headers.py
213-213: Local variable graph is assigned to but never used
Remove assignment to unused variable graph
(F841)
🔇 Additional comments (8)
app.py (4)
30-30: Flask import update is appropriateAdding
has_request_contextto the Flask imports is correct and required by the new isolation helpers.
2484-2489: store_memory isolation wiring looks correct for the synchronous pathWithin
store_memory:
- You resolve
graph_name = _get_graph_name()andcollection_name = _get_collection_name().- Use
graph = get_memory_graph(graph_name)for the FalkorDB write.- Use
collection_namefor the synchronous Qdrant upsert when an embedding is provided.Aside from the async embedding concern noted separately, this synchronous path correctly honors the per-request isolation headers.
Also applies to: 2570-2593
2645-2649: update_memory correctly applies per-request isolation to graph and QdrantFor
update_memory:
- You resolve
graph_name/collection_namevia the new helpers.- Use
get_memory_graph(graph_name)for the graph update.- Use
collection_nameconsistently for Qdrantretrieveandupsert.This ensures updates are scoped to the requested graph/collection. Existing memories stored in the default collection will still be reachable when called without headers.
Also applies to: 2738-2763
78-82: CI reports Black reformatting this fileGitHub Actions indicates that
blackreformattedapp.py. Please run:
black app.pyisort app.py --profile=blacklocally (or
black ./isort .) to align with the repository’s formatting and unbreak CI.As per coding guidelines, Python files should be Black- and isort-formatted.
tests/test_isolation_headers.py (4)
7-13: Imports and module setup look fine, but ensure formatting passes CIThe imports and top-level setup are standard for a pytest module targeting
app. Since CI reports that Black/isort reformatted this file, make sure to re-run:black tests/test_isolation_headers.py isort tests/test_isolation_headers.py --profile=blackso the committed file matches what CI expects.
As per coding guidelines, test files should also be Black- and isort-formatted.
35-141: Header extraction and validation tests are thoroughThe tests in
TestHeaderExtractionandTestHeaderValidationcover:
- Valid graph/collection names (including underscores and hyphens).
- Defaulting behavior when headers are missing or whitespace.
- Regex/length constraints and HTTP 400 responses for invalid values.
They align well with the logic in
_get_graph_name/_get_collection_nameand should catch regressions in header parsing.
144-201: Whitelist behavior tests correctly exercise ALLOWED_GRAPHS/COLLECTIONSThe
TestWhitelistEnforcementcases verify:
- Accepting listed names and rejecting others with 403.
- Handling whitespace in env-configured lists.
- Behavior when no whitelist is set (any valid name allowed).
This matches the intended semantics of the isolation whitelists.
222-271: Connection caching tests correctly validate per-name behaviorThe
TestConnectionCachingtests:
- Patch
app.stateto a MagicMock to avoid real FalkorDB connections.- Clear
_graph_cachebefore each scenario.- Assert that repeated calls with the same graph name reuse the cached instance and don’t re-call
select_graph.- Assert that different graph names result in separate instances and multiple
select_graphinvocations.This is a solid verification of the new
_graph_cachebehavior.
| class TestBackwardsCompatibility: | ||
| """Test that existing behavior is preserved when headers are not provided.""" | ||
|
|
||
| def test_get_memory_graph_without_request_context_uses_default(self): | ||
| """Test that get_memory_graph() works outside request context.""" | ||
| # This should not raise an error and should use the default | ||
| # Note: This will fail if FalkorDB is not available, but that's expected | ||
| # in a unit test environment. The key is that it doesn't raise a RuntimeError | ||
| # about missing request context. | ||
| try: | ||
| graph = app.get_memory_graph() | ||
| # If it returns None, that's fine - FalkorDB might not be running | ||
| # We're just testing it doesn't crash | ||
| except RuntimeError as e: | ||
| if "request context" in str(e).lower(): | ||
| pytest.fail("get_memory_graph() should not require request context") | ||
| # Other RuntimeErrors are acceptable (e.g., connection failures) | ||
|
|
There was a problem hiding this comment.
Fix unused local graph in backwards-compatibility test
In TestBackwardsCompatibility.test_get_memory_graph_without_request_context_uses_default:
graph = app.get_memory_graph()graph is never used, and Ruff flags this as F841.
You can just call the function without binding the result:
- try:
- graph = app.get_memory_graph()
+ try:
+ app.get_memory_graph()This still validates that no RuntimeError about missing request context is raised, without introducing an unused local.
🧰 Tools
🪛 Ruff (0.14.8)
213-213: Local variable graph is assigned to but never used
Remove assignment to unused variable graph
(F841)
🤖 Prompt for AI Agents
In tests/test_isolation_headers.py around lines 203 to 220, the
backwards-compatibility test assigns the result of app.get_memory_graph() to an
unused local variable `graph`, which triggers a Ruff F841 warning; remove the
assignment and call app.get_memory_graph() directly (i.e., replace `graph =
app.get_memory_graph()` with `app.get_memory_graph()`) so the function is
exercised and any RuntimeError about request context would be caught without
creating an unused variable.
1 participant
Summary
Adds per-request isolation headers (
X-Graph-Name,X-Collection-Name) to enable multi-tenant usage of a single AutoMem instance. This is a minimal, backwards-compatible change that allows federation layers to route memories to separate graph/collection namespaces.Changes
_get_graph_name()and_get_collection_name()extract and validate isolation headersALLOWED_GRAPHSandALLOWED_COLLECTIONSenv vars for securitystore_memory,update_memory,delete_memory,recall_memoriesto use isolation headershas_request_context()check ensures consolidation/enrichment use defaultsBackwards Compatibility
GRAPH_NAMEandCOLLECTION_NAMEenv defaults)Use Case
This enables the automem-federation layer to route memories to different banks using a single AutoMem deployment with multiple FalkorDB graphs and Qdrant collections.
Test Plan