Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions engraphis/mcp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
7 changes: 6 additions & 1 deletion engraphis/routes/v2_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
58 changes: 58 additions & 0 deletions engraphis/service_context.py
Original file line number Diff line number Diff line change
@@ -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)
70 changes: 70 additions & 0 deletions tests/test_service_context.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
"""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():
with require_service_context(), pytest.raises(RuntimeError, match="context is required"):
bound_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):
principal = {"id": "member_%d" % index, "email": "u%d@example.test" % index,
"role": "member"}
with bind_service(services[index], principal=principal):
bound_service().remember("Only tenant %d" % index, workspace="shared")
barrier.wait(timeout=5)
assert bound_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_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")
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()