diff --git a/integrations/crew-ai/python/README.md b/integrations/crew-ai/python/README.md index ffefd4fbf..4b70e62e6 100644 --- a/integrations/crew-ai/python/README.md +++ b/integrations/crew-ai/python/README.md @@ -78,6 +78,56 @@ opens, and a `RAW` first event makes the reference client reject the whole strea those are held and released once the run has opened. Both RAW buffers are bounded; a saturated buffer degrades mirroring (logged) rather than the run. +### Memory is isolated per `threadId` (default ON) + +A crew served with `Crew(memory=True)` keeps its memories in one on-disk store, +namespaced by the *crew name*. Nothing in that namespace derives from the AG-UI +`threadId`, so without help every chat served by an endpoint reads and writes the +same namespace and one user's remembered facts surface in another user's chat. +(Setting `inputs["id"] = thread_id` does not help: that scopes crewai's flow-state +persistence, a different subsystem.) `Agent(memory=True)` has the same shape one +level down: the agent builds its *own* memory, which crewai prefers over the +crew's. + +The bridge closes that by giving each request a `MemoryScope` view of the crew's +memory — and of each agent's own memory — rooted at a path derived from the +request's `threadId`. Threads are mutually invisible; each still sees its own +history across sequential runs. One physical store, no directory-per-thread +sprawl. + +Because crewai picks the executing agent off `task.agent` (or `manager_agent` +under the hierarchical process) and reaches the crew's memory through +`agent.crew`, the request gets shallow *views* of the crew, its agents and its +tasks, wired to each other. Nothing shared between concurrent requests is +mutated, and everything below the views (tools, LLMs, knowledge, the store +itself) stays shared. + +```bash +# Opt out: restore the pre-fix behaviour of one memory namespace per crew, +# shared by every chat. Useful when the crew is a durable knowledge base +# rather than a per-conversation memory. +AGUI_CREWAI_THREAD_SCOPED_MEMORY=false +``` + +Limitations, in order of how likely you are to hit them: + +- **Only crews and agents the bridge can reach are scoped.** That means the crew + you passed to `add_crewai_crew_fastapi_endpoint`, plus any crew or standalone + agent your `Flow` holds as an attribute (a crew's own agents and tasks come + with it). A crew or agent *constructed inside* a flow method is created after + this point and is not scoped; construct it as a flow attribute, or pass it a + `Memory` you scope yourself. +- **Per-request views are shallow.** Each request runs against copies of the + crew, its agents and its tasks, so `crew.tasks[0].output` on the object you + built is not filled in by a bridge-served run; read the run's result off the + AG-UI event stream instead. +- **Isolation is logical, not physical.** All threads share one store and one + embedder; a scope keeps reads and writes inside a namespace, it is not a + security boundary against code that queries the store directly. +- **Older crewai degrades rather than crashing.** The bridge probes for crewai's + unified memory view API at runtime. On a build without it, isolation is not + active and the bridge logs one warning saying exactly that. + ### `get_capabilities()` Returns a capability declaration (the CrewAI counterpart of diff --git a/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py b/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py index 431ab9493..9e3f624f6 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py +++ b/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py @@ -276,6 +276,14 @@ def flow_supports_stream_frames(flow: Any) -> bool: _Flow = getattr(_CREWAI_MODULE, "Flow", None) if _CREWAI_MODULE else None _Crew = getattr(_CREWAI_MODULE, "Crew", None) if _CREWAI_MODULE else None +# ``BaseAgent`` is the base every crewai agent derives from, including a user's +# own subclass, so it is the wider net for "this attribute is an agent". +# ``crewai.Agent`` is the fallback for a build that does not expose it. +_BASE_AGENT_MODULE, _ = _first_module(["crewai.agents.agent_builder.base_agent"]) +_Agent = ( + getattr(_BASE_AGENT_MODULE, "BaseAgent", None) if _BASE_AGENT_MODULE else None +) or (getattr(_CREWAI_MODULE, "Agent", None) if _CREWAI_MODULE else None) + def _kwarg_in_signature(func: Any, name: str) -> bool: """True when ``func`` declares a parameter ``name`` (or accepts ``**kwargs``). @@ -388,6 +396,46 @@ def supported_checkpoint_kwargs(method: Any, kwargs: dict[str, Any]) -> dict[str return {k: v for k, v in kwargs.items() if k in params} +# -------------------------------------------------------------------------- +# Memory-isolation resolution +# -------------------------------------------------------------------------- +# crewai 1.x replaced the 0.x short-term / entity / long-term stores with ONE +# ``Memory`` object over ONE store, namespaced by a ``root_scope`` string that +# ``Crew.create_crew_memory`` derives from the CREW NAME. Nothing in that path +# derives from the AG-UI ``threadId``, so two chats served by the same endpoint +# read and write the same namespace. +# +# The isolation primitive is ``Memory.scope(path)``, which returns a +# ``MemoryScope`` view whose reads and writes are confined to ``path`` and +# below. ``Crew._memory`` is typed ``Memory | MemoryScope | MemorySlice``, so a +# view is a first-class thing to hand a crew, not a hack. +# +# Resolved (never version-gated) so a build without the unified memory API +# degrades to "no isolation, one warning" rather than crashing. The warning is +# emitted at the CALL SITE (``_memory``) rather than from ``warn_on_gaps``: an +# operator who never sets ``memory=True`` has no gap to hear about, and an +# import-time warning for them would be pure noise. +_MEMORY_MODULE, _MEMORY_MODULE_NAME = _first_module(["crewai.memory.unified_memory"]) +Memory = getattr(_MEMORY_MODULE, "Memory", None) if _MEMORY_MODULE is not None else None + +# crewai's own scope-name sanitizer (``crewai.memory.utils``). Used so a +# bridge-built scope segment is normalised exactly the way crewai normalises the +# crew-name segment sitting above it. ``_memory`` carries an equivalent fallback +# for builds that do not expose it. +_MEMORY_UTILS_MODULE, _ = _first_module(["crewai.memory.utils"]) +sanitize_scope_name = ( + getattr(_MEMORY_UTILS_MODULE, "sanitize_scope_name", None) + if _MEMORY_UTILS_MODULE is not None + else None +) + +# Both are required: the type (to recognise a crew's memory) and the view +# factory (to derive a per-thread namespace from it). +_memory_scope_available = Memory is not None and callable( + getattr(Memory, "scope", None) +) + + @dataclass(frozen=True) class _Capabilities: """Cached, immutable snapshot of the detected crewai capabilities. @@ -415,6 +463,12 @@ class _Capabilities: checkpoint_fork_supported: bool = False checkpoint_events_available: bool = False checkpoint_state_module: str | None = None + # Per-thread memory isolation: informational. ``_memory`` keys off the + # resolved symbols and warns once at the call site, so this is NOT listed in + # ``missing`` (which drives import-time warnings that would fire for every + # operator, including the majority who never enable crew memory). + memory_scope_available: bool = False + memory_module: str | None = None missing: tuple[str, ...] = field(default_factory=tuple) def warn_on_gaps(self) -> None: @@ -488,6 +542,8 @@ def _detect() -> _Capabilities: checkpoint_fork_supported=_checkpoint_fork_supported, checkpoint_events_available=_checkpoint_events_available, checkpoint_state_module=_CKPT_STATE_MODULE_NAME, + memory_scope_available=_memory_scope_available, + memory_module=_MEMORY_MODULE_NAME, missing=tuple(missing), ) caps.warn_on_gaps() diff --git a/integrations/crew-ai/python/ag_ui_crewai/_config.py b/integrations/crew-ai/python/ag_ui_crewai/_config.py index 664b51698..afefd3d3a 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/_config.py +++ b/integrations/crew-ai/python/ag_ui_crewai/_config.py @@ -26,6 +26,15 @@ EMIT_RAW_EVENTS_ENV_VAR = "AGUI_CREWAI_EMIT_RAW_EVENTS" +# Per-thread memory isolation (crew and agent) ships ON: sharing one namespace across +# every AG-UI ``threadId`` leaks one chat's remembered facts into another, which +# is a privacy bug rather than a feature. The opt-out exists because a +# deployment may WANT one durable knowledge base behind every chat; turning it +# off restores the pre-fix "one namespace per crew name" behaviour exactly. +DEFAULT_THREAD_SCOPED_MEMORY = True + +THREAD_SCOPED_MEMORY_ENV_VAR = "AGUI_CREWAI_THREAD_SCOPED_MEMORY" + # Vocabulary ``_parse_env_bool`` accepts, so the "was this value used?" check stays # in step with the parser instead of duplicating its token list. _BOOL_TOKENS = _TRUE_VALUES | _FALSE_VALUES @@ -76,3 +85,28 @@ def resolve_emit_raw_events(emit_raw_events: bool | None) -> bool: used = raw is not None and raw.strip().casefold() in _BOOL_TOKENS _warn_if_env_value_ignored(EMIT_RAW_EVENTS_ENV_VAR, raw, used) return resolved + + +def resolve_thread_scoped_memory() -> bool: + """Resolve per-thread crew-memory isolation: env var > shipped default (on). + + Env-only, and re-read per request rather than resolved once at registration: + unlike ``emit_raw_events`` there is no endpoint-factory argument to conflict + with, and an operator flipping the variable should not have to know which + call it was frozen at. + + Deliberately NOT ``_parse_env_bool``: that parser treats anything outside its + true-set as false, which is the right fail-safe for an option that ships OFF + but the wrong one here: a typo would silently DISABLE isolation and restore + the cross-thread leak. Only a recognised false token turns it off; anything + else keeps the shipped default and warns once. + """ + raw = os.environ.get(THREAD_SCOPED_MEMORY_ENV_VAR) + if raw is None: + return DEFAULT_THREAD_SCOPED_MEMORY + token = raw.strip().casefold() + used = token in _BOOL_TOKENS + _warn_if_env_value_ignored(THREAD_SCOPED_MEMORY_ENV_VAR, raw, used) + if not used: + return DEFAULT_THREAD_SCOPED_MEMORY + return token in _TRUE_VALUES diff --git a/integrations/crew-ai/python/ag_ui_crewai/_memory.py b/integrations/crew-ai/python/ag_ui_crewai/_memory.py new file mode 100644 index 000000000..a49196c12 --- /dev/null +++ b/integrations/crew-ai/python/ag_ui_crewai/_memory.py @@ -0,0 +1,423 @@ +"""Per-thread isolation of CrewAI memory. + +The bug. A crew served with ``Crew(memory=True)`` leaks remembered +facts between chats. The bridge isolates requests by copying the flow, but that +copy does not touch memory: crewai 1.x builds ONE ``Memory`` over ONE on-disk +store, namespaced by a ``root_scope`` string that ``Crew.create_crew_memory`` +derives from the CREW NAME. Every request to a given endpoint therefore reads +and writes the same namespace. Passing ``inputs["id"] = thread_id`` does not +help; that scopes crewai's flow-state persistence, a different subsystem. +``Agent(memory=True)`` has the same shape one level down: the agent builds its +OWN ``Memory``, and the executor prefers it over the crew's +(``memory = agent.memory or crew._memory`` in ``BaseAgentExecutor``), so an +agent-level memory bypasses a crew-level fix entirely. + +The fix. Before the run starts, give THIS request's flow copy a crew whose +memory -- and whose agents' memories -- are ``MemoryScope`` views rooted at a +path derived from the AG-UI ``threadId``. Reads and writes through such a view +are confined to the thread's namespace and below, so thread A and thread B are +mutually invisible while each still sees its own history across sequential runs. +One physical store, no directory sprawl, and both ``Crew._memory`` and +``Agent.memory`` are typed ``Memory | MemoryScope | MemorySlice`` upstream, so a +view is a supported thing to hand them rather than a hack. + +Three constraints shape the implementation: + +* **The template crew is shared across concurrent requests, and so are its + agents and tasks.** ``ChatWithCrewFlow`` is built once per endpoint and cached, + and ``_copyutil.safe_deepcopy`` pins the uncopyable crew BY REFERENCE, so every + in-flight request holds the same live ``Crew``, ``Agent``s and ``Task``s. + Re-pointing any of those shared objects per request would be a data race. So + nothing shared is ever mutated: we build per-request shallow VIEWS (crew, + agents, tasks), point only the views at the thread scope, and swap the crew + view onto the per-request flow copy. ``_verify_template_untouched`` fails the + request loudly if a future refactor ever writes through to a shared object. +* **The memory has to reach the object that actually executes.** crewai resolves + the executing agent from ``task.agent`` (or ``Crew.manager_agent`` under the + hierarchical process), NOT from ``Crew.agents``, and it reaches the crew's + memory through ``agent.crew`` -- which ``Crew.kickoff`` ASSIGNS on each agent + it is given. Scoping an agent view that no task points at would therefore + scope nothing, and leaving the shared agents in place would let one request's + ``agent.crew = `` decide which thread's namespace another request + reads. So the whole executed graph is re-pointed together: agent views on the + crew view, and task views whose ``agent`` is the matching agent view. +* **There are two run paths.** The legacy ``kickoff_async`` driver and the + ``astream`` StreamFrame driver both read the crew off the flow copy, so scoping + the flow copy before either driver is selected covers both. + +Capability-detected, never version-gated: a crewai build without the unified +``Memory.scope`` view API degrades to "isolation not active" plus one warning +rather than crashing. +""" + +from __future__ import annotations + +import copy +import hashlib +import logging +import re +from typing import Any + +# ``_Crew`` / ``_Agent`` / ``sanitize_scope_name`` are resolved in +# ``_capabilities`` with every other crewai symbol the bridge depends on, rather +# than re-derived here with a parallel ``getattr`` chain. +from ._capabilities import ( + CAPABILITIES, + _Agent, + _Crew, + sanitize_scope_name as _crewai_sanitize_scope_name, +) +from ._config import resolve_thread_scoped_memory + +_LOGGER = logging.getLogger(__name__) + +# Namespace the per-thread scopes live under, relative to whatever root scope the +# crew already carries. A crew named "support" therefore stores thread A's +# memories under ``/crew/support/thread/`` and keeps crewai's own +# hierarchy intact above it. +_THREAD_SCOPE_ROOT = "/thread" + +# Upper bound on the readable part of a scope segment. Thread ids are usually +# UUIDs, but nothing stops a client sending a very long one, and the segment ends +# up in a stored scope path. +_MAX_READABLE_SEGMENT = 64 + +# Length of the digest suffix, in hex characters. 16 hex chars is 64 bits, which +# makes an accidental collision between two live threads unreachable in practice. +_DIGEST_CHARS = 16 + +# Emitted at most once per process: an operator whose crewai lacks the view API +# needs to hear that isolation is off, but not once per request. +_DEGRADE_WARNED = False + +# Raised when a write meant for a per-request view lands on the shared template +# instead. That is strictly worse than the bug this module fixes -- it would hand +# every other in-flight request THIS thread's namespace -- so it fails the run +# rather than degrading quietly. +_SHARED_MUTATION = ( + "ag-ui-crewai: building the per-thread memory views mutated the SHARED " + "{what} instead of the per-request copy. Refusing to continue; concurrent " + "requests would read one another's memory." +) + + +def _fallback_sanitize_scope_name(name: str) -> str: + """Stand-in for ``crewai.memory.utils.sanitize_scope_name``. + + Same normalisation (lowercase, non ``[a-z0-9_-]`` to a hyphen, collapse and + strip hyphens) so a build that does not expose crewai's helper still produces + a segment shaped like the crew-name segment above it. + """ + if not name: + return "unknown" + name = re.sub(r"[^a-z0-9_-]", "-", name.lower().strip()) + name = re.sub(r"-+", "-", name).strip("-") + return name or "unknown" + + +def thread_scope_path(thread_id: str) -> str: + """Return the scope path that isolates ``thread_id``. + + Shape: ``/thread/-``. The readable half is the sanitized + thread id, so a stored scope tree stays diagnosable by eye. The digest half + is what actually guarantees isolation: sanitisation is lossy (``a/b`` and + ``a-b`` both normalise to ``a-b``, and a long id is truncated), so two + distinct threads could otherwise collapse onto one namespace and leak into + each other, which is the exact bug this module exists to prevent. The digest + is taken over the RAW id and is therefore injective in practice. + """ + raw = str(thread_id) + sanitize = _crewai_sanitize_scope_name or _fallback_sanitize_scope_name + try: + readable = sanitize(raw) + except Exception: # noqa: BLE001 - a surprising crewai helper must not fail the run + readable = _fallback_sanitize_scope_name(raw) + readable = readable[:_MAX_READABLE_SEGMENT].strip("-") or "unknown" + digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:_DIGEST_CHARS] + return f"{_THREAD_SCOPE_ROOT}/{readable}-{digest}" + + +def _warn_isolation_unavailable(memory: Any, what: str) -> None: + """Say once, loudly, that this deployment is running without isolation. + + Two causes worth distinguishing: the installed crewai has no unified memory + view API at all, or it does but this particular crew (or agent) was handed a + memory object that is not one of crewai's view types. The remedy differs. + """ + global _DEGRADE_WARNED # pylint: disable=global-statement + if _DEGRADE_WARNED: + return + _DEGRADE_WARNED = True + if not CAPABILITIES.memory_scope_available: + cause = ( + f"crewai {CAPABILITIES.crewai_version} does not expose the unified " + "memory view API (crewai.memory.unified_memory.Memory.scope)" + ) + remedy = "Upgrade crewai, or disable memory" + else: + cause = ( + f"this {what}'s memory is a {type(memory).__name__}, which has no " + "scope() view factory" + ) + remedy = ( + f"Pass the {what} a crewai Memory (or an already-scoped view), or " + "disable memory" + ) + _LOGGER.warning( + "ag-ui-crewai: %s memory is enabled but %s, so PER-THREAD MEMORY " + "ISOLATION IS NOT ACTIVE: every AG-UI threadId served by this endpoint " + "shares one memory namespace and can read another chat's remembered " + "facts. %s. Further occurrences are silenced.", + what, + cause, + remedy, + ) + + +def _attribute_containers(obj: Any) -> list[dict]: + """Return the mutable attribute stores of ``obj`` worth scanning. + + Pydantic keeps declared fields in ``__dict__`` and private attributes in + ``__pydantic_private__``; a crewai ``Flow`` holds an assigned ``self.crew`` in + the former. Writing back through these dicts (rather than ``setattr``) is the + same approach ``_copyutil.rebind_bound_methods`` takes, and for the same + reason: ``BaseModel.__setattr__`` refuses or reroutes assignments that are not + declared fields. + """ + containers = [] + for candidate in ( + getattr(obj, "__dict__", None), + getattr(obj, "__pydantic_private__", None), + ): + if isinstance(candidate, dict): + containers.append(candidate) + return containers + + +def _write_attr(obj: Any, name: str, value: Any) -> None: + """Point ``obj.name`` at ``value``, going through the attribute store. + + Pydantic keeps declared fields (``Agent.memory``, ``Task.agent``, + ``Crew.agents``) in ``__dict__`` and private attributes (``Crew._memory``) in + ``__pydantic_private__``; writing back through whichever dict already holds + the name avoids ``BaseModel.__setattr__`` refusing or re-validating the + assignment. Only ever called on a ``copy.copy`` VIEW, which owns fresh + mappings, so the shared original is unaffected. + """ + private = getattr(obj, "__pydantic_private__", None) + if isinstance(private, dict) and name in private: + private[name] = value + return + fields = getattr(obj, "__dict__", None) + if isinstance(fields, dict) and name in fields: + fields[name] = value + return + setattr(obj, name, value) + + +def _has_memory(owner: Any) -> bool: + """True when ``owner`` carries a memory object of its own. + + ``True`` / ``False`` do not count: crewai's validators turn those into a + ``Memory`` / ``None`` at construction, so a leftover bool is a + not-yet-validated value, not a memory to scope. + """ + memory = getattr(owner, "memory", None) + return memory is not None and not isinstance(memory, bool) + + +def _scoped_memory(memory: Any, scope_path: str, what: str) -> Any | None: + """Return a ``scope_path``-rooted view of ``memory``, or ``None``. + + ``None`` covers all three "nothing to do here" cases: no memory, a memory + object this crewai build cannot make a view of (warned about once), and a + view factory that raised (warned about per occurrence, because a crewai-side + error is worth seeing every time it happens). + """ + if memory is None or isinstance(memory, bool): + return None + + scope_factory = getattr(memory, "scope", None) + if not callable(scope_factory): + _warn_isolation_unavailable(memory, what) + return None + + try: + return scope_factory(scope_path) + except Exception as exc: # noqa: BLE001 - never fail a run over memory scoping + _LOGGER.warning( + "ag-ui-crewai: could not scope %s memory to %s (%s: %s); this run " + "shares that memory's namespace with other threads.", + what, + scope_path, + type(exc).__name__, + exc, + ) + return None + + +def _thread_scoped_agent(agent: Any, scope_path: str) -> Any: + """Return a per-request view of ``agent``, its own memory scoped if it has one. + + Always a view, even for an agent with no memory of its own: ``Crew.kickoff`` + ASSIGNS ``agent.crew`` (and ``agent.agent_executor``) on every agent it runs, + and the executor reads the crew's memory back off that attribute. Handing two + concurrent requests the same agent object would let the later one decide + which thread's namespace the earlier one reads. + """ + memory = getattr(agent, "memory", None) + view = copy.copy(agent) + scoped_memory = _scoped_memory(memory, scope_path, "agent") + if scoped_memory is not None: + _write_attr(view, "memory", scoped_memory) + if getattr(agent, "memory", None) is not memory: + raise RuntimeError(_SHARED_MUTATION.format(what="Agent.memory")) + return view + + +def _thread_scoped_tasks(tasks: Any, agent_views: dict[int, Any]) -> list[Any] | None: + """Return per-request task views pointing at ``agent_views``, or ``None``. + + crewai picks the executing agent off ``task.agent``, so a task left pointing + at the shared agent would bypass the scoped view entirely. ``task.context`` + is remapped alongside, because it names OTHER task objects whose ``output`` + this run fills in -- left pointing at the shared tasks, a downstream task + would read an output this request never produced. + """ + if not isinstance(tasks, (list, tuple)): + return None + + views = [copy.copy(task) for task in tasks] + by_task = {id(task): view for task, view in zip(tasks, views)} + for task, view in zip(tasks, views): + agent_view = agent_views.get(id(getattr(task, "agent", None))) + if agent_view is not None: + _write_attr(view, "agent", agent_view) + context = getattr(task, "context", None) + if isinstance(context, list) and any(id(t) in by_task for t in context): + _write_attr(view, "context", [by_task.get(id(t), t) for t in context]) + return views + + +def _verify_template_untouched(crew: Any, before: dict[tuple[int, str], Any]) -> None: + """Fail loudly if building the views wrote through to a shared object. + + ``before`` maps ``(id(shared object), attribute)`` to the value it held + before the views were built. A mismatch means one of the writes landed on the + template rather than on its copy, which under concurrency hands every other + in-flight request THIS thread's namespace -- a worse leak than the one this + module exists to fix, and one that would otherwise be silent. + """ + for owner, attr in _template_attrs(crew): + key = (id(owner), attr) + if key in before and getattr(owner, attr, None) is not before[key]: + raise RuntimeError( + _SHARED_MUTATION.format(what=f"{type(owner).__name__}.{attr}") + ) + + +def _template_attrs(crew: Any): + """Yield every ``(shared object, attribute)`` pair the view build must not touch.""" + yield crew, "_memory" + yield crew, "agents" + yield crew, "tasks" + yield crew, "manager_agent" + for agent in _crew_agents(crew): + yield agent, "memory" + yield agent, "crew" + tasks = getattr(crew, "tasks", None) + if isinstance(tasks, (list, tuple)): + for task in tasks: + yield task, "agent" + yield task, "context" + + +def _crew_agents(crew: Any) -> list[Any]: + """Every agent the crew can execute with: its roster plus the manager.""" + agents = getattr(crew, "agents", None) + roster = list(agents) if isinstance(agents, (list, tuple)) else [] + manager = getattr(crew, "manager_agent", None) + if manager is not None: + roster.append(manager) + return roster + + +def _thread_scoped_crew(crew: Any, scope_path: str) -> Any | None: + """Return a shallow view of ``crew`` whose memory is scoped to ``scope_path``. + + ``None`` means "leave this crew alone": neither the crew nor any of its + agents has memory to scope, or this crewai build cannot produce a view. + + The passed crew, its agents and its tasks are SHARED with every concurrent + request and are never mutated -- ``copy.copy`` of a Pydantic model builds + fresh ``__dict__`` / ``__pydantic_private__`` mappings, so pointing a view's + attribute elsewhere leaves the original's alone. Everything below the views + (tools, LLMs, knowledge, the underlying store, the save pool) stays shared, + exactly as it already was before this module existed. + """ + scoped_memory = _scoped_memory(getattr(crew, "_memory", None), scope_path, "crew") + agents = _crew_agents(crew) + if scoped_memory is None and not any(_has_memory(agent) for agent in agents): + return None + + before = {(id(o), a): getattr(o, a, None) for o, a in _template_attrs(crew)} + + agent_views = {id(agent): _thread_scoped_agent(agent, scope_path) for agent in agents} + + crew_view = copy.copy(crew) + if scoped_memory is not None: + _write_attr(crew_view, "_memory", scoped_memory) + roster = getattr(crew, "agents", None) + if isinstance(roster, (list, tuple)): + _write_attr(crew_view, "agents", [agent_views[id(a)] for a in roster]) + manager = getattr(crew, "manager_agent", None) + if manager is not None: + _write_attr(crew_view, "manager_agent", agent_views[id(manager)]) + task_views = _thread_scoped_tasks(getattr(crew, "tasks", None), agent_views) + if task_views is not None: + _write_attr(crew_view, "tasks", task_views) + + _verify_template_untouched(crew, before) + return crew_view + + +def apply_thread_memory_scope(flow_copy: Any, thread_id: str | None) -> None: + """Scope every crew and agent on ``flow_copy`` to ``thread_id``'s namespace. + + Called once per request, on the PER-REQUEST flow copy, before a run driver is + selected -- so both the legacy ``kickoff_async`` path and the ``astream`` + StreamFrame path are covered. + + Scans the flow copy's own attributes for ``Crew`` instances (whose agents and + tasks are scoped with them) and for standalone ``Agent`` instances, which a + flow can drive directly through ``agent.kickoff()`` without a crew. That + reaches the crew-serving endpoint's ``ChatWithCrewFlow.crew`` and whatever a + user's ``Flow`` holds as an attribute. A crew or agent CONSTRUCTED INSIDE a + flow method is not reachable from here and is not scoped; see the README for + that limitation. + + Never raises for an environment reason: an unscopable memory degrades to the + previous shared-namespace behaviour with a warning, because failing a chat + outright is worse than the leak it would prevent. + """ + if not resolve_thread_scoped_memory(): + return + if not thread_id: + # Nothing to derive a namespace from. Scoping every such request onto one + # shared "unknown" namespace would isolate nothing while still hiding the + # crew's own history, so leave the crew as-is. + return + if _Crew is None: # pragma: no cover - crewai is a hard dependency + return + + scope_path = thread_scope_path(thread_id) + for container in _attribute_containers(flow_copy): + for name, value in list(container.items()): + if isinstance(value, _Crew): + crew_view = _thread_scoped_crew(value, scope_path) + if crew_view is not None: + container[name] = crew_view + elif _Agent is not None and isinstance(value, _Agent) and _has_memory(value): + agent_view = _thread_scoped_agent(value, scope_path) + if getattr(agent_view, "memory", None) is not value.memory: + container[name] = agent_view diff --git a/integrations/crew-ai/python/ag_ui_crewai/endpoint.py b/integrations/crew-ai/python/ag_ui_crewai/endpoint.py index 3b6042038..46602793b 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/endpoint.py +++ b/integrations/crew-ai/python/ag_ui_crewai/endpoint.py @@ -34,6 +34,7 @@ from ._checkpoint import build_checkpoint_kwargs from ._frames import StreamFrameTranslator, CREW_AGENT_LIFECYCLE_TYPES from ._config import resolve_emit_raw_events as _resolve_emit_raw_events +from ._memory import apply_thread_memory_scope from ._frames import is_recognized_event, log_raw_loss, raw_event_for from .mcp import is_mcp_event, register_mcp_listeners, translate_mcp_event from .attribution import flat_method_attribution @@ -2262,6 +2263,11 @@ async def agentic_chat_endpoint(input_data: RunAgentInput, request: Request): """Agentic chat endpoint""" flow_copy = _copy_flow(flow) + # Copying the flow does NOT isolate crew memory: crewai's memory lives in + # a shared on-disk store namespaced by crew name, so every threadId would + # otherwise read and write the same namespace. Scope it here, before a run + # driver is selected, so both drivers are covered. + apply_thread_memory_scope(flow_copy, input_data.thread_id) # Get the accept header from the request accept_header = request.headers.get("accept") @@ -2348,6 +2354,11 @@ async def crew_endpoint(input_data: RunAgentInput, request: Request): """Crew chat endpoint with deferred initialization.""" flow = await _get_flow() flow_copy = _copy_flow(flow) + # See the note in ``add_crewai_flow_fastapi_endpoint``. This is the path + # the reported ``Crew(memory=True)`` leak was filed against: the cached + # ``ChatWithCrewFlow`` holds one crew, pinned by reference into every + # request's copy, so the scope view is what separates the threads. + apply_thread_memory_scope(flow_copy, input_data.thread_id) accept_header = request.headers.get("accept") encoder = EventEncoder(accept=accept_header) diff --git a/integrations/crew-ai/python/tests/conftest.py b/integrations/crew-ai/python/tests/conftest.py index 0dca2f624..e6ddaa096 100644 --- a/integrations/crew-ai/python/tests/conftest.py +++ b/integrations/crew-ai/python/tests/conftest.py @@ -12,16 +12,45 @@ """ import copy +import os +import shutil +import tempfile import pytest -from ag_ui_crewai import endpoint as ep +# Redirect crewai's on-disk storage root BEFORE anything imports crewai. +# +# crewai resolves its storage root at MODULE-IMPORT time (``crewai.rag.chromadb. +# constants`` calls ``db_storage_path()``, which ``mkdir(parents=True)``s the +# directory), so merely importing the bridge writes into the developer's home +# directory and a per-test fixture would already be too late. The root comes from +# ``CREWAI_STORAGE_DIR``, which MUST be absolute: ``db_storage_path`` treats a +# relative value as an appdirs *app name* (landing back in ``$HOME``) while +# ``LanceDBStorage`` treats it as a *directory* (landing in the cwd). +# +# This only makes the TEST RUN hermetic. The import-time write itself is crewai's +# behaviour, not the bridge's, and cannot be suppressed from library code without +# setting environment variables on the user's behalf. +_OWNED_STORAGE_DIR = None +if not os.environ.get("CREWAI_STORAGE_DIR"): + _OWNED_STORAGE_DIR = tempfile.mkdtemp(prefix="ag-ui-crewai-tests-") + os.environ["CREWAI_STORAGE_DIR"] = _OWNED_STORAGE_DIR + +from ag_ui_crewai import endpoint as ep # noqa: E402 # The crewai global event bus — used below to clear handlers registered by our # listener singleton so they don't accumulate across tests. # The bus moved from ``crewai.utilities.events`` (0.x) to # ``crewai.events`` (1.x); ``_capabilities`` resolves whichever exists. -from ag_ui_crewai._capabilities import crewai_event_bus as _crewai_event_bus +from ag_ui_crewai._capabilities import crewai_event_bus as _crewai_event_bus # noqa: E402 + + +@pytest.fixture(scope="session", autouse=True) +def _cleanup_crewai_storage_dir(): + """Remove the temporary crewai storage root this session created, if any.""" + yield + if _OWNED_STORAGE_DIR: + shutil.rmtree(_OWNED_STORAGE_DIR, ignore_errors=True) # crewai 1.0.0 split the single ``_handlers`` mapping into # ``_sync_handlers`` / ``_async_handlers``. The autouse fixture below snapshots diff --git a/integrations/crew-ai/python/tests/test_memory_isolation.py b/integrations/crew-ai/python/tests/test_memory_isolation.py new file mode 100644 index 000000000..5fa1a7b91 --- /dev/null +++ b/integrations/crew-ai/python/tests/test_memory_isolation.py @@ -0,0 +1,710 @@ +"""Per-thread isolation of CrewAI memory (crew-level and agent-level). + +Everything here runs OFFLINE against a real ``Crew(memory=True)`` / a real +``Agent`` with its own ``Memory``: a fake embedding function stands in for +OpenAI, and every write supplies an explicit scope / categories / importance, +which is the condition under which crewai's ``EncodingFlow`` skips its LLM +analysis step. The only thing mocked is the embedder, so the store, the scope +arithmetic and the recall path are all real. + +``Agent(memory=True)`` builds its Memory with crewai's DEFAULT embedder, which +needs a live API key the moment anything is written. Behavioural tests therefore +hand the agent an explicit ``Memory(embedder=)`` -- the same +object ``memory=True`` would have produced, minus the network. Constructing +``Agent(memory=True)`` is itself offline (the embedder is resolved lazily), so +the reported shape is still asserted directly, structurally. + +The seam under test is the one the FastAPI endpoints use: copy the flow for this +request, then scope the copy's memory to the request's ``threadId``. +""" + +import copy +import hashlib +import logging + +import pytest +from crewai import Agent, Crew, Flow, Process, Task +from crewai.flow import start +from crewai.memory.unified_memory import Memory +from chromadb.api.types import EmbeddingFunction as _ChromaEmbeddingFunction +from crewai.rag.embeddings.providers.custom.embedding_callable import ( + CustomEmbeddingFunction as _CrewEmbeddingFunction, +) + +from ag_ui_crewai import _memory as memory_module +from ag_ui_crewai._memory import apply_thread_memory_scope, thread_scope_path +from ag_ui_crewai.endpoint import _copy_flow + + +class _DeterministicEmbedder(_CrewEmbeddingFunction, _ChromaEmbeddingFunction): + """Offline embedder: the first 16 bytes of a SHA-256 digest, scaled to [0, 1]. + + crewai validates the ``custom`` embedder spec against chromadb's + ``EmbeddingFunction`` and then against its own ``CustomEmbeddingFunction``, + so the class has to satisfy both. + """ + + def __init__(self): + # chromadb warns when an embedding function omits ``__init__``. + pass + + def __call__(self, input): # noqa: A002 - the upstream parameter is named ``input`` + return [ + [byte / 255.0 for byte in hashlib.sha256(text.encode()).digest()[:16]] + for text in input + ] + + +_EMBEDDER_SPEC = { + "provider": "custom", + "config": {"embedding_callable": _DeterministicEmbedder}, +} + + +def _remember(view, content): + """Write through a ``Memory`` or ``MemoryScope`` without invoking an LLM.""" + return view.remember( + content, scope="/facts", categories=["fact"], importance=0.9 + ) + + +def _recall(view, query): + return [match.record.content for match in view.recall(query, depth="shallow")] + + +@pytest.fixture +def memory_crew(tmp_path, monkeypatch): + """A real ``Crew(memory=True)`` whose store lives under ``tmp_path``. + + ``CREWAI_STORAGE_DIR`` is read when the storage backend is CONSTRUCTED, which + happens inside ``Crew(...)`` via the ``create_crew_memory`` model validator, + so setting it here (before the crew is built) is early enough. It must be + absolute; see the note in ``conftest.py``. + """ + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + agent = Agent(role="helper", goal="help", backstory="b", llm="gpt-4o-mini") + task = Task(description="d", expected_output="e", agent=agent) + return Crew( + name="support-crew", + agents=[agent], + tasks=[task], + memory=True, + embedder=_EMBEDDER_SPEC, + ) + + +class _CrewHoldingFlow(Flow): + """Minimal stand-in for ``ChatWithCrewFlow``: a flow with a ``crew`` attribute. + + ``ChatWithCrewFlow.__init__`` issues a live LLM call to generate its chat + inputs, so it cannot be constructed offline. The only property this seam cares + about is the one reproduced here: a crew reachable as a flow attribute, shared + by every request because the flow is cached per endpoint. + """ + + def __init__(self, crew): + super().__init__() + self.crew = crew + + @start() + def go(self): # pragma: no cover - never kicked off in these tests + return "ok" + + +@pytest.fixture +def crew_flow(memory_crew): + return _CrewHoldingFlow(memory_crew) + + +def _serve(flow, thread_id): + """Do to ``flow`` exactly what an endpoint does at the start of a request.""" + flow_copy = _copy_flow(flow) + apply_thread_memory_scope(flow_copy, thread_id) + return flow_copy + + +@pytest.fixture(autouse=True) +def _reset_degrade_warning(): + """Un-latch the one-shot degradation warning between tests.""" + memory_module._DEGRADE_WARNED = False + yield + memory_module._DEGRADE_WARNED = False + + +# --------------------------------------------------------------------------- +# The reported bug +# --------------------------------------------------------------------------- + + +def test_two_threads_cannot_read_each_others_memory(crew_flow): + """The reported leak: chat B must not see what chat A remembered.""" + run_a = _serve(crew_flow, "thread-A") + _remember(run_a.crew._memory, "THREAD-A-SECRET: the user's name is Ada") + + run_b = _serve(crew_flow, "thread-B") + + assert _recall(run_b.crew._memory, "what is the user's name") == [] + # ...and A did not lose its own memory in the process. + assert any( + "THREAD-A-SECRET" in content + for content in _recall(run_a.crew._memory, "what is the user's name") + ) + + +def test_writes_from_both_threads_stay_separated(crew_flow): + """Symmetry: neither direction leaks once both threads have written.""" + run_a = _serve(crew_flow, "thread-A") + run_b = _serve(crew_flow, "thread-B") + + _remember(run_a.crew._memory, "A-ONLY: the favourite colour is teal") + _remember(run_b.crew._memory, "B-ONLY: the favourite colour is amber") + + assert _recall(run_a.crew._memory, "favourite colour") == [ + "A-ONLY: the favourite colour is teal" + ] + assert _recall(run_b.crew._memory, "favourite colour") == [ + "B-ONLY: the favourite colour is amber" + ] + + +def test_one_thread_keeps_its_memory_across_sequential_runs(crew_flow): + """Isolation must not degenerate into amnesia: a thread still has a history.""" + first_run = _serve(crew_flow, "thread-A") + _remember(first_run.crew._memory, "REMEMBERED: the deploy target is staging") + + second_run = _serve(crew_flow, "thread-A") + third_run = _serve(crew_flow, "thread-A") + + assert _recall(second_run.crew._memory, "deploy target") == [ + "REMEMBERED: the deploy target is staging" + ] + assert _recall(third_run.crew._memory, "deploy target") == [ + "REMEMBERED: the deploy target is staging" + ] + + +# --------------------------------------------------------------------------- +# Race safety: the template crew is shared across concurrent requests +# --------------------------------------------------------------------------- + + +def test_shared_template_crew_is_never_mutated(crew_flow): + """Scoping must not re-point the crew every other in-flight request holds.""" + template_crew = crew_flow.crew + template_memory = template_crew._memory + + run = _serve(crew_flow, "thread-A") + + assert crew_flow.crew is template_crew + assert template_crew._memory is template_memory + assert type(template_memory).__name__ == "Memory" + # The request got its own crew view carrying a scoped memory. + assert run.crew is not template_crew + assert type(run.crew._memory).__name__ == "MemoryScope" + + +def test_concurrent_requests_get_independent_scopes(crew_flow): + """Two overlapping requests must not clobber one another's scope.""" + run_a = _serve(crew_flow, "thread-A") + run_b = _serve(crew_flow, "thread-B") + + assert run_a.crew is not run_b.crew + assert run_a.crew._memory.root_path != run_b.crew._memory.root_path + # Both views still share the one physical store; only the namespace differs. + assert run_a.crew._memory._memory is run_b.crew._memory._memory + + +def test_scoped_view_changes_nothing_but_the_memory(crew_flow): + """The crew, agent and task views differ from the originals only in memory. + + The agents and tasks are per-request views (see + ``test_the_shared_agents_and_tasks_are_never_mutated`` for why), but they are + still equal to the originals field-for-field, and everything hanging off them + -- tools, LLM, knowledge -- is the same object. + """ + run = _serve(crew_flow, "thread-A") + + assert run.crew.name == crew_flow.crew.name + assert run.crew.agents == crew_flow.crew.agents + assert run.crew.tasks == crew_flow.crew.tasks + assert run.crew.agents[0].llm is crew_flow.crew.agents[0].llm + assert run.crew.agents[0].tools is crew_flow.crew.agents[0].tools + + +# --------------------------------------------------------------------------- +# Agent-level memory: ``Agent(memory=...)`` beats the crew's, so it needs its own +# scope. crewai's executor resolves ``agent.memory or crew._memory``. +# --------------------------------------------------------------------------- + + +def _offline_memory(): + """The ``Memory`` ``Agent(memory=True)`` builds, minus the OpenAI embedder.""" + return Memory(embedder=_EMBEDDER_SPEC) + + +@pytest.fixture +def agent_memory_flow(tmp_path, monkeypatch): + """A crew whose AGENT carries its own memory (the crew's is off).""" + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + agent = Agent( + role="helper", + goal="help", + backstory="b", + llm="gpt-4o-mini", + memory=_offline_memory(), + ) + task = Task(description="d", expected_output="e", agent=agent) + crew = Crew(name="support-crew", agents=[agent], tasks=[task]) + return _CrewHoldingFlow(crew) + + +def _executing_agent(run): + """The agent crewai will actually run the first task with.""" + return run.crew._get_agent_to_use(run.crew.tasks[0]) + + +def test_two_threads_cannot_read_each_others_agent_memory(agent_memory_flow): + """The same leak one level down: an agent's own memory must be per-thread.""" + run_a = _serve(agent_memory_flow, "thread-A") + _remember(_executing_agent(run_a).memory, "AGENT-A-SECRET: the user is Ada") + + run_b = _serve(agent_memory_flow, "thread-B") + + assert _recall(_executing_agent(run_b).memory, "who is the user") == [] + assert _recall(_executing_agent(run_a).memory, "who is the user") == [ + "AGENT-A-SECRET: the user is Ada" + ] + + +def test_one_thread_keeps_its_agent_memory_across_sequential_runs(agent_memory_flow): + """Isolation must not degenerate into amnesia at the agent level either.""" + first_run = _serve(agent_memory_flow, "thread-A") + _remember(_executing_agent(first_run).memory, "REMEMBERED: the target is staging") + + second_run = _serve(agent_memory_flow, "thread-A") + + assert _recall(_executing_agent(second_run).memory, "the target") == [ + "REMEMBERED: the target is staging" + ] + + +def test_the_task_executes_with_the_scoped_agent(agent_memory_flow): + """The scoped agent has to be the one crewai picks, or the scoping is inert. + + crewai resolves the executing agent from ``task.agent``, NOT from + ``Crew.agents``, so a task left pointing at the shared agent would run with + the shared, unscoped memory however well-scoped the roster is. + """ + run = _serve(agent_memory_flow, "thread-A") + + assert _executing_agent(run) is run.crew.agents[0] + assert type(_executing_agent(run).memory).__name__ == "MemoryScope" + + +def test_an_agent_built_with_memory_true_is_scoped(tmp_path, monkeypatch): + """The reported shape, asserted verbatim: ``Agent(memory=True)``. + + Structural only -- ``memory=True`` resolves to a ``Memory`` on crewai's + DEFAULT embedder, so writing through it would need a live API key. + """ + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + agent = Agent(role="helper", goal="help", backstory="b", llm="gpt-4o-mini", memory=True) + task = Task(description="d", expected_output="e", agent=agent) + flow = _CrewHoldingFlow(Crew(name="support-crew", agents=[agent], tasks=[task])) + + run_a = _serve(flow, "thread-A") + run_b = _serve(flow, "thread-B") + + assert type(agent.memory).__name__ == "Memory" + assert type(_executing_agent(run_a).memory).__name__ == "MemoryScope" + assert ( + _executing_agent(run_a).memory.root_path + != _executing_agent(run_b).memory.root_path + ) + + +def test_a_hierarchical_manager_agents_memory_is_scoped(tmp_path, monkeypatch): + """Under the hierarchical process the executing agent is the MANAGER.""" + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + manager = Agent( + role="manager", + goal="manage", + backstory="b", + llm="gpt-4o-mini", + memory=_offline_memory(), + ) + worker = Agent(role="helper", goal="help", backstory="b", llm="gpt-4o-mini") + task = Task(description="d", expected_output="e", agent=worker) + flow = _CrewHoldingFlow( + Crew( + name="support-crew", + agents=[worker], + tasks=[task], + process=Process.hierarchical, + manager_agent=manager, + ) + ) + + run_a = _serve(flow, "thread-A") + _remember(_executing_agent(run_a).memory, "MANAGER-A: the budget is 40") + + run_b = _serve(flow, "thread-B") + + assert _executing_agent(run_a) is run_a.crew.manager_agent + assert _recall(_executing_agent(run_b).memory, "the budget") == [] + assert type(manager.memory).__name__ == "Memory" + + +def test_a_standalone_agent_on_the_flow_is_scoped(tmp_path, monkeypatch): + """An agent a flow drives directly, with no crew, is scoped too.""" + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + agent = Agent( + role="helper", + goal="help", + backstory="b", + llm="gpt-4o-mini", + memory=_offline_memory(), + ) + + class _AgentHoldingFlow(Flow): + def __init__(self): + super().__init__() + self.agent = agent + + @start() + def go(self): # pragma: no cover - never kicked off + return "ok" + + flow = _AgentHoldingFlow() + run_a = _serve(flow, "thread-A") + _remember(run_a.agent.memory, "SOLO-A: the room is 12B") + + run_b = _serve(flow, "thread-B") + + assert _recall(run_b.agent.memory, "which room") == [] + assert type(agent.memory).__name__ == "Memory" + + +def test_an_agent_without_its_own_memory_still_gets_the_scoped_crew(crew_flow): + """Falling back to crew memory must reach THIS request's scoped crew. + + ``Crew.kickoff`` assigns ``agent.crew`` on every agent it is given, and the + executor reads the crew's memory back off that attribute. Two concurrent + requests sharing one agent object would leave the loser reading the winner's + thread namespace, so each request gets its own agent view even when the agent + has no memory of its own. + """ + run_a = _serve(crew_flow, "thread-A") + run_b = _serve(crew_flow, "thread-B") + + assert _executing_agent(run_a).memory is None + assert _executing_agent(run_a) is not _executing_agent(run_b) + assert _executing_agent(run_a) is not crew_flow.crew.agents[0] + + +def test_the_shared_agents_and_tasks_are_never_mutated(agent_memory_flow): + """Nothing the concurrent requests share may be re-pointed.""" + template_crew = agent_memory_flow.crew + template_agent = template_crew.agents[0] + template_task = template_crew.tasks[0] + template_memory = template_agent.memory + + run_a = _serve(agent_memory_flow, "thread-A") + run_b = _serve(agent_memory_flow, "thread-B") + + assert template_agent.memory is template_memory + assert type(template_memory).__name__ == "Memory" + assert template_task.agent is template_agent + assert template_crew.agents[0] is template_agent + assert template_crew.tasks[0] is template_task + # ...and the two requests hold independent scopes over the one store. + memory_a = _executing_agent(run_a).memory + memory_b = _executing_agent(run_b).memory + assert memory_a.root_path != memory_b.root_path + assert memory_a._memory is memory_b._memory is template_memory + + +def test_task_context_edges_point_at_this_requests_tasks(tmp_path, monkeypatch): + """A copied task's ``context`` must name the copies, not the shared tasks. + + ``context`` names other Task OBJECTS, and crewai reads their ``.output`` to + build the downstream prompt. Left pointing at the shared tasks, a downstream + task would read an output this request never produced. + """ + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + agent = Agent(role="helper", goal="help", backstory="b", llm="gpt-4o-mini") + first = Task(description="first", expected_output="e", agent=agent) + second = Task(description="second", expected_output="e", agent=agent, context=[first]) + flow = _CrewHoldingFlow( + Crew( + name="support-crew", + agents=[agent], + tasks=[first, second], + memory=True, + embedder=_EMBEDDER_SPEC, + ) + ) + + run = _serve(flow, "thread-A") + + assert run.crew.tasks[1].context == [run.crew.tasks[0]] + assert run.crew.tasks[1].context[0] is run.crew.tasks[0] + assert second.context == [first] + + +def test_opt_out_leaves_agent_memory_shared(agent_memory_flow, monkeypatch): + """``AGUI_CREWAI_THREAD_SCOPED_MEMORY=false`` disables agent scoping too.""" + monkeypatch.setenv("AGUI_CREWAI_THREAD_SCOPED_MEMORY", "false") + + run_a = _serve(agent_memory_flow, "thread-A") + _remember(_executing_agent(run_a).memory, "SHARED: the office is in Lisbon") + run_b = _serve(agent_memory_flow, "thread-B") + + assert _executing_agent(run_a) is agent_memory_flow.crew.agents[0] + assert _recall(_executing_agent(run_b).memory, "where is the office") == [ + "SHARED: the office is in Lisbon" + ] + + +# --------------------------------------------------------------------------- +# Fail-loud: a write that lands on a shared object must not be served +# --------------------------------------------------------------------------- + + +def test_a_write_through_to_the_shared_agent_fails_the_request(tmp_path, monkeypatch): + """A copy that is not a copy must raise, not silently share one namespace.""" + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + + class _UncopyableAgent(Agent): + def __copy__(self): + return self + + agent = _UncopyableAgent( + role="helper", + goal="help", + backstory="b", + llm="gpt-4o-mini", + memory=_offline_memory(), + ) + task = Task(description="d", expected_output="e", agent=agent) + flow = _CrewHoldingFlow(Crew(name="support-crew", agents=[agent], tasks=[task])) + + with pytest.raises(RuntimeError, match="mutated the SHARED .*Agent.memory"): + _serve(flow, "thread-A") + + +def test_a_write_through_to_the_shared_task_fails_the_request(tmp_path, monkeypatch): + """Same guarantee for tasks: ``task.agent`` is what selects the memory.""" + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + + class _UncopyableTask(Task): + def __copy__(self): + return self + + agent = Agent(role="helper", goal="help", backstory="b", llm="gpt-4o-mini") + task = _UncopyableTask(description="d", expected_output="e", agent=agent) + flow = _CrewHoldingFlow( + Crew( + name="support-crew", + agents=[agent], + tasks=[task], + memory=True, + embedder=_EMBEDDER_SPEC, + ) + ) + + with pytest.raises(RuntimeError, match="mutated the SHARED .*Task.agent"): + _serve(flow, "thread-A") + + +# --------------------------------------------------------------------------- +# Scope-path derivation +# --------------------------------------------------------------------------- + + +def test_scope_path_is_derived_from_the_thread_id(): + path = thread_scope_path("Thread A") + assert path.startswith("/thread/thread-a-") + assert thread_scope_path("Thread A") == path + + +def test_scope_paths_differ_for_ids_that_sanitise_alike(): + """Sanitisation is lossy; the digest suffix is what keeps threads apart.""" + assert thread_scope_path("a/b") != thread_scope_path("a-b") + assert thread_scope_path("x" * 200) != thread_scope_path("x" * 201) + + +def test_scope_path_survives_an_id_with_no_usable_characters(): + assert thread_scope_path("///").startswith("/thread/unknown-") + + +def test_two_threads_whose_ids_sanitise_alike_are_isolated(crew_flow): + """The end-to-end consequence of the digest suffix.""" + run_slash = _serve(crew_flow, "tenant/7") + run_dash = _serve(crew_flow, "tenant-7") + + _remember(run_slash.crew._memory, "SLASH-TENANT: quota is 40") + + assert _recall(run_dash.crew._memory, "quota") == [] + + +# --------------------------------------------------------------------------- +# Opt-out +# --------------------------------------------------------------------------- + + +def test_opt_out_restores_the_shared_namespace(crew_flow, monkeypatch): + """``AGUI_CREWAI_THREAD_SCOPED_MEMORY=false`` reinstates one namespace per crew.""" + monkeypatch.setenv("AGUI_CREWAI_THREAD_SCOPED_MEMORY", "false") + + run_a = _serve(crew_flow, "thread-A") + _remember(run_a.crew._memory, "SHARED: the office is in Lisbon") + run_b = _serve(crew_flow, "thread-B") + + assert run_a.crew is crew_flow.crew + assert _recall(run_b.crew._memory, "where is the office") == [ + "SHARED: the office is in Lisbon" + ] + + +def test_an_unrecognised_opt_out_value_keeps_isolation_on(crew_flow, monkeypatch): + """A typo must not silently reopen the leak.""" + monkeypatch.setenv("AGUI_CREWAI_THREAD_SCOPED_MEMORY", "flase") + + run_a = _serve(crew_flow, "thread-A") + _remember(run_a.crew._memory, "TYPO-GUARD: the office is in Lisbon") + run_b = _serve(crew_flow, "thread-B") + + assert _recall(run_b.crew._memory, "where is the office") == [] + + +# --------------------------------------------------------------------------- +# Degradation and no-ops +# --------------------------------------------------------------------------- + + +def test_a_crew_without_memory_is_left_untouched(tmp_path, monkeypatch): + monkeypatch.setenv("CREWAI_STORAGE_DIR", str(tmp_path / "crewai-store")) + agent = Agent(role="helper", goal="help", backstory="b", llm="gpt-4o-mini") + task = Task(description="d", expected_output="e", agent=agent) + crew = Crew(name="memoryless", agents=[agent], tasks=[task]) + flow = _CrewHoldingFlow(crew) + + run = _serve(flow, "thread-A") + + assert run.crew._memory is None + + +def test_a_missing_thread_id_leaves_the_crew_untouched(crew_flow): + run = _serve(crew_flow, "") + + assert type(run.crew._memory).__name__ == "Memory" + + +def test_a_memory_without_the_view_api_degrades_with_one_warning(crew_flow, caplog): + """No crash, no isolation, and exactly one warning per process.""" + + class _LegacyMemory: + """A memory object predating (or replacing) ``Memory.scope``.""" + + unscopable_crew = copy.copy(crew_flow.crew) + unscopable_crew.__pydantic_private__["_memory"] = _LegacyMemory() + flow = _CrewHoldingFlow(unscopable_crew) + + with caplog.at_level(logging.WARNING, logger="ag_ui_crewai._memory"): + first = _serve(flow, "thread-A") + second = _serve(flow, "thread-B") + + assert isinstance(first.crew._memory, _LegacyMemory) + assert isinstance(second.crew._memory, _LegacyMemory) + warnings = [ + record for record in caplog.records + if "PER-THREAD MEMORY ISOLATION IS NOT ACTIVE" in record.message + ] + assert len(warnings) == 1 + + +def test_a_failing_scope_factory_degrades_instead_of_failing_the_run( + crew_flow, caplog +): + """A crewai-side error while building the view must not kill the chat.""" + + class _AngryMemory: + def scope(self, path): + raise RuntimeError("boom") + + angry_crew = copy.copy(crew_flow.crew) + angry_crew.__pydantic_private__["_memory"] = _AngryMemory() + flow = _CrewHoldingFlow(angry_crew) + + with caplog.at_level(logging.WARNING, logger="ag_ui_crewai._memory"): + run = _serve(flow, "thread-A") + + assert isinstance(run.crew._memory, _AngryMemory) + assert any("could not scope crew memory" in r.message for r in caplog.records) + + +# --------------------------------------------------------------------------- +# Endpoint wiring +# --------------------------------------------------------------------------- + + +class _InertFlow: + """A flow-shaped object that neither runs nor supports StreamFrames.""" + + async def kickoff_async(self, inputs=None): # pragma: no cover - never awaited + return "ok" + + +@pytest.mark.parametrize("factory", ["flow", "crew"]) +async def test_both_endpoint_factories_scope_memory_before_running( + factory, monkeypatch +): + """Both endpoints scope the request's flow copy, and do it BEFORE a driver runs. + + Scoping happens in the endpoint body, ahead of ``_run_flow_stream``'s choice + between the legacy ``kickoff_async`` driver and the ``astream`` StreamFrame + driver, so it cannot be dead code on one of them. + """ + from types import SimpleNamespace + + from ag_ui.core import RunAgentInput + from fastapi import FastAPI + + from ag_ui_crewai import endpoint as ep + + calls = [] + monkeypatch.setattr( + ep, + "apply_thread_memory_scope", + lambda flow_copy, thread_id: calls.append((flow_copy, thread_id)), + ) + + app = FastAPI() + if factory == "flow": + ep.add_crewai_flow_fastapi_endpoint(app, _InertFlow(), path="/run") + else: + monkeypatch.setattr( + ep, "ChatWithCrewFlow", lambda *_a, **_kw: _InertFlow() + ) + ep.add_crewai_crew_fastapi_endpoint(app, object(), path="/run") + + route = next(r for r in app.router.routes if getattr(r, "path", None) == "/run") + request = SimpleNamespace(headers=SimpleNamespace(get=lambda *_a, **_kw: None)) + await route.endpoint( + RunAgentInput( + thread_id="thread-A", + run_id="run-1", + state={}, + messages=[], + tools=[], + context=[], + forwarded_props={}, + ), + request, + ) + + assert len(calls) == 1 + scoped_flow, thread_id = calls[0] + assert thread_id == "thread-A" + assert isinstance(scoped_flow, _InertFlow)