fix(mcp): scope active workspace per caller identity

This commit is contained in:
CREDO23 2026-07-07 16:52:30 +02:00
parent 93bd022c02
commit 44d0f28c5f

View file

@ -7,14 +7,22 @@ speaks a name, we resolve it, and remember the choice for later calls.
from __future__ import annotations from __future__ import annotations
from collections import OrderedDict
from dataclasses import dataclass from dataclasses import dataclass
from typing import Annotated from typing import Annotated
from pydantic import Field from pydantic import Field
from .auth.identity import current_identity
from .client import SurfSenseClient from .client import SurfSenseClient
from .errors import ToolError from .errors import ToolError
# ponytail: one small entry per distinct caller (API token). Bounded so a flood
# of keys can't grow memory without limit; an evicted caller just re-resolves
# its default workspace on the next call. Upgrade path: a TTL/LRU store if
# per-caller state ever grows past this one field.
_MAX_TRACKED_IDENTITIES = 2048
# Shared parameter type for every workspace-scoped tool. # Shared parameter type for every workspace-scoped tool.
WorkspaceParam = Annotated[ WorkspaceParam = Annotated[
str | None, str | None,
@ -44,15 +52,21 @@ class WorkspaceContext:
) -> None: ) -> None:
self._client = client self._client = client
self._preferred_reference = preferred_reference self._preferred_reference = preferred_reference
self._active: Workspace | None = None # Active selection is per caller: one shared slot would leak one user's
# choice to every other user on a shared server.
self._active_by_identity: OrderedDict[str, Workspace] = OrderedDict()
@property @property
def active(self) -> Workspace | None: def active(self) -> Workspace | None:
return self._active return self._active_by_identity.get(current_identity())
def remember(self, workspace: Workspace) -> Workspace: def remember(self, workspace: Workspace) -> Workspace:
"""Make ``workspace`` the default for later scoped calls.""" """Make ``workspace`` the default for the current caller's later calls."""
self._active = workspace identity = current_identity()
self._active_by_identity[identity] = workspace
self._active_by_identity.move_to_end(identity)
while len(self._active_by_identity) > _MAX_TRACKED_IDENTITIES:
self._active_by_identity.popitem(last=False)
return workspace return workspace
async def fetch_all(self) -> list[Workspace]: async def fetch_all(self) -> list[Workspace]:
@ -67,8 +81,9 @@ class WorkspaceContext:
return self.remember(await self._match(reference)) return self.remember(await self._match(reference))
async def _resolve_default(self) -> Workspace: async def _resolve_default(self) -> Workspace:
if self._active is not None: active = self.active
return self._active if active is not None:
return active
if self._preferred_reference: if self._preferred_reference:
return self.remember(await self._match(self._preferred_reference)) return self.remember(await self._match(self._preferred_reference))
return self.remember(await self._only_workspace_or_prompt()) return self.remember(await self._only_workspace_or_prompt())