From bb122d1e6f21c7c04100ebad3d420f3fb032ab4a Mon Sep 17 00:00:00 2001 From: Coding-Dev-Tools Date: Sun, 6 Sep 2026 01:31:20 -0400 Subject: [PATCH 1/2] feat: support request-scoped memory service contexts --- engraphis/mcp_server.py | 5 +++ engraphis/routes/v2_api.py | 7 ++++- engraphis/service_context.py | 58 ++++++++++++++++++++++++++++++++++ tests/test_service_context.py | 59 +++++++++++++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 1 deletion(-) create mode 100644 engraphis/service_context.py create mode 100644 tests/test_service_context.py diff --git a/engraphis/mcp_server.py b/engraphis/mcp_server.py index 784836e6..b811c312 100644 --- a/engraphis/mcp_server.py +++ b/engraphis/mcp_server.py @@ -111,6 +111,11 @@ def set_service(svc: MemoryService) -> None: def service() -> MemoryService: """Lazily build the service so server startup is instant (model loads on first use).""" + from engraphis.service_context import bound_service + + bound = bound_service() + if bound is not None: + return bound global _service if _service is None: with _service_lock: diff --git a/engraphis/routes/v2_api.py b/engraphis/routes/v2_api.py index 4b810e54..f5353259 100644 --- a/engraphis/routes/v2_api.py +++ b/engraphis/routes/v2_api.py @@ -115,7 +115,12 @@ def _sanitized_http_exception(status_code: object) -> HTTPException: def service() -> MemoryService: - """Lazily bind a single MemoryService to the configured store (the live v2 DB).""" + """Resolve the request binding or lazily open the standalone local store.""" + from engraphis.service_context import bound_service + + bound = bound_service() + if bound is not None: + return bound global _service with _SERVICE_LOCK: if _service is None: diff --git a/engraphis/service_context.py b/engraphis/service_context.py new file mode 100644 index 00000000..be7f90f6 --- /dev/null +++ b/engraphis/service_context.py @@ -0,0 +1,58 @@ +"""Request-scoped service injection for applications embedding the local engine. + +The standalone entry points retain their local default. A hosted application must +enter ``bind_service`` for each operation, with a validated principal. Contexts +are restored even when an operation fails, and never mutate module singletons. +""" +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from typing import TYPE_CHECKING, Iterator, Optional + +if TYPE_CHECKING: + from engraphis.service import MemoryService + +_BOUND: ContextVar[Optional["MemoryService"]] = ContextVar("engraphis_bound_service", default=None) +_REQUIRED: ContextVar[bool] = ContextVar("engraphis_bound_service_required", default=False) + + +def bound_service() -> Optional["MemoryService"]: + """Resolve an injected service, refusing an absent required binding.""" + result = _BOUND.get() + if result is None and _REQUIRED.get(): + raise RuntimeError("An authenticated service context is required") + return result + + +@contextmanager +def require_service_context() -> Iterator[None]: + """Disable the standalone fallback for the duration of a hosted request.""" + token = _REQUIRED.set(True) + try: + yield + finally: + _REQUIRED.reset(token) + + +@contextmanager +def bind_service(service: "MemoryService", *, principal: dict) -> Iterator["MemoryService"]: + """Bind an explicit service and principal, restoring the enclosing context.""" + from engraphis.service import _CURRENT_USER, set_current_user + + if service is None or not principal: + raise ValueError("An explicit service and authenticated principal are required") + previous_user = _CURRENT_USER.get() + try: + set_current_user(principal) + except Exception: + _CURRENT_USER.set(previous_user) + raise + service_token = _BOUND.set(service) + required_token = _REQUIRED.set(True) + try: + yield service + finally: + _REQUIRED.reset(required_token) + _BOUND.reset(service_token) + _CURRENT_USER.set(previous_user) diff --git a/tests/test_service_context.py b/tests/test_service_context.py new file mode 100644 index 00000000..8be8065f --- /dev/null +++ b/tests/test_service_context.py @@ -0,0 +1,59 @@ +"""Concurrent embedding contexts must never inherit another tenant's service.""" +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier + +import pytest + +from engraphis.service import MemoryService, current_user +from engraphis.service_context import bind_service, bound_service, require_service_context + + +def test_required_context_cannot_fall_back_to_local_service(): + from engraphis.routes.v2_api import service + + with require_service_context(), pytest.raises(RuntimeError, match="context is required"): + service() + assert bound_service() is None + + +def test_contexts_keep_identical_workspace_names_isolated(): + services = [MemoryService.create(":memory:", extractor="none") for _ in range(2)] + barrier = Barrier(2) + + def work(index): + from engraphis.routes.v2_api import service + + principal = {"id": "member_%d" % index, "email": "u%d@example.test" % index, + "role": "member"} + with bind_service(services[index], principal=principal): + service().remember("Only tenant %d" % index, workspace="shared") + barrier.wait(timeout=5) + assert service() is services[index] + assert current_user()["id"] == principal["id"] + assert bound_service() is None + assert current_user() is None + + try: + with ThreadPoolExecutor(max_workers=2) as pool: + list(pool.map(work, range(2))) + finally: + for instance in services: + instance.close() + + +def test_exception_and_nested_binding_restore_outer_identity(): + first = MemoryService.create(":memory:", extractor="none") + second = MemoryService.create(":memory:", extractor="none") + user = {"id": "member_outer", "email": "outer@example.test", "role": "member"} + other = {"id": "member_inner", "email": "inner@example.test", "role": "viewer"} + try: + with bind_service(first, principal=user): + with pytest.raises(ValueError, match="operation failed"): + with bind_service(second, principal=other): + raise ValueError("operation failed") + assert bound_service() is first + assert current_user()["id"] == user["id"] + assert current_user() is None + finally: + first.close() + second.close() From 401f787555de2d9384f8573012ca179d12e2f44a Mon Sep 17 00:00:00 2001 From: Coding-Dev-Tools Date: Sun, 6 Sep 2026 02:04:01 -0400 Subject: [PATCH 2/2] Test service contexts independently of optional HTTP dependencies --- tests/test_service_context.py | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/tests/test_service_context.py b/tests/test_service_context.py index 8be8065f..c6371d80 100644 --- a/tests/test_service_context.py +++ b/tests/test_service_context.py @@ -9,10 +9,8 @@ def test_required_context_cannot_fall_back_to_local_service(): - from engraphis.routes.v2_api import service - with require_service_context(), pytest.raises(RuntimeError, match="context is required"): - service() + bound_service() assert bound_service() is None @@ -21,14 +19,12 @@ def test_contexts_keep_identical_workspace_names_isolated(): barrier = Barrier(2) def work(index): - from engraphis.routes.v2_api import service - principal = {"id": "member_%d" % index, "email": "u%d@example.test" % index, "role": "member"} with bind_service(services[index], principal=principal): - service().remember("Only tenant %d" % index, workspace="shared") + bound_service().remember("Only tenant %d" % index, workspace="shared") barrier.wait(timeout=5) - assert service() is services[index] + assert bound_service() is services[index] assert current_user()["id"] == principal["id"] assert bound_service() is None assert current_user() is None @@ -41,6 +37,21 @@ def work(index): instance.close() +def test_http_adapter_requires_explicit_context_when_requested(): + pytest.importorskip("fastapi", reason="HTTP adapter requires the optional server extra") + from engraphis.routes.v2_api import service + + with require_service_context(), pytest.raises(RuntimeError, match="context is required"): + service() + instance = MemoryService.create(":memory:", extractor="none") + principal = {"id": "member_http", "email": "http@example.test", "role": "member"} + try: + with bind_service(instance, principal=principal): + assert service() is instance + finally: + instance.close() + + def test_exception_and_nested_binding_restore_outer_identity(): first = MemoryService.create(":memory:", extractor="none") second = MemoryService.create(":memory:", extractor="none")