Python: scope FIDES quarantine clients to their owning config

python/packages/core/agent_framework/security.py

@@ -1379,6 +1379,8 @@ def __init__(
         self.auto_hide_untrusted = auto_hide_untrusted
         self.hide_threshold = hide_threshold
         self._standing_guidance_cache: WeakKeyDictionary[Any, tuple[str, ...]] = WeakKeyDictionary()
+        self._quarantine_client_bound = False
+        self._quarantine_client_override: SupportsChatGetResponse | None = None
         self._initialize_security_scope(security_scope, session_state_key=session_state_key)
 
     def _clone_for_scope(self, scope: _SecurityScope) -> LabelTrackingFunctionMiddleware:
@@ -3364,15 +3366,11 @@ class SecureAgentConfig(ContextProvider):
         auto_hide_untrusted: Whether to automatically hide untrusted content.
 
     Note:
-        The quarantine chat client is stored per-instance (see ``get_quarantine_client``)
-        but is *also* registered in a single process-global slot via
-        ``set_quarantine_client``. The ``quarantined_llm`` tool always reads that global
-        slot, so the behavior is **last-writer-wins**: when multiple ``SecureAgentConfig``
-        instances are constructed in the same process with different ``quarantine_chat_client``
-        values, the most recently constructed instance's client is the one every agent's
-        ``quarantined_llm`` tool will use. Running multiple instances is supported, but they
-        share this one global quarantine client rather than each using their own. If you need
-        distinct quarantine clients per agent, run them in separate processes.
+        During calls mediated by this config's label-tracking middleware,
+        ``quarantined_llm`` uses this config's ``quarantine_chat_client``.
+        An explicitly missing client cannot inherit one from another config.
+        Direct standalone calls without an active config retain the legacy
+        process-global fallback configured by ``set_quarantine_client``.
 
     Examples:
         .. code-block:: python
@@ -3433,9 +3431,8 @@ def __init__(
                 If provided, the quarantined_llm tool will make actual isolated LLM calls
                 instead of returning placeholder responses. This client should ideally be
                 a separate instance using a cheaper model (e.g., gpt-4o-mini) since it
-                processes untrusted content. Note: this client is registered in a
-                process-global slot shared by all instances (last-writer-wins); see the
-                class docstring for details on running multiple instances.
+                processes untrusted content. Scoped agent calls use this client;
+                standalone calls retain the legacy process-global fallback.
             source_id: Optional source identifier for context provider attribution.
                 Defaults to "secure_agent".
 
@@ -3492,14 +3489,23 @@ def _scope_for_session(self, session: AgentSession) -> _SecurityScope:
             raise ValueError("Security provider state must be a dictionary.")
         return _SecurityScope(cast(dict[str, Any], provider_state), scope_id=session.session_id)
 
+    def _bind_quarantine_client_to_stack(self, middleware_stack: list[FunctionMiddleware]) -> list[FunctionMiddleware]:
+        """Bind this config's client to each active tracker, including replacements."""
+        for middleware in middleware_stack:
+            if isinstance(middleware, LabelTrackingFunctionMiddleware):
+                middleware._quarantine_client_bound = True  # pyright: ignore[reportPrivateUsage]
+                middleware._quarantine_client_override = self._quarantine_chat_client  # pyright: ignore[reportPrivateUsage]
+        return middleware_stack
+
     def _middleware_for_scope(self, scope: _SecurityScope) -> list[FunctionMiddleware]:
         """Clone the current overridable middleware stack into one scope."""
-        return [
+        scoped = [
             middleware._clone_for_scope(scope)  # pyright: ignore[reportPrivateUsage]
             if isinstance(middleware, (LabelTrackingFunctionMiddleware, PolicyEnforcementFunctionMiddleware))
             else middleware
             for middleware in self.get_middleware()
         ]
+        return self._bind_quarantine_client_to_stack(scoped)
 
     def _function_middleware_for_approval_resolution(self, session: AgentSession) -> list[FunctionMiddleware]:
         """Return provider-customized middleware bound to one restored session scope."""
@@ -3557,7 +3563,7 @@ def get_middleware(self) -> list[FunctionMiddleware]:
         middleware: list[FunctionMiddleware] = [self.label_tracker]
         if self.policy_enforcer:
             middleware.append(self.policy_enforcer)
-        return middleware
+        return self._bind_quarantine_client_to_stack(middleware)
 
     def get_audit_log(self, session: AgentSession | None = None) -> list[dict[str, Any]]:
         """Get the audit log for an optional session.
@@ -3974,7 +3980,12 @@ async def quarantined_llm(
     actual_metadata = metadata or {}
 
     # Build the response - use real LLM if quarantine client is configured
-    quarantine_client = get_quarantine_client()
+    # Scoped security middleware wins, including an explicitly missing client.
+    quarantine_client = (
+        middleware._quarantine_client_override  # pyright: ignore[reportPrivateUsage]
+        if middleware is not None and middleware._quarantine_client_bound  # pyright: ignore[reportPrivateUsage]
+        else get_quarantine_client()
+    )
 
     if quarantine_client is not None:
         # Build the quarantined prompt with retrieved content

python/packages/core/tests/test_security.py

@@ -24,6 +24,7 @@
     FunctionMiddleware,
     Message,
     SessionContext,
+    SupportsChatGetResponse,
 )
 from agent_framework._middleware import FunctionMiddlewarePipeline, MiddlewareFailure, MiddlewareTermination
 from agent_framework._tools import (
@@ -8841,3 +8842,121 @@ async def my_tool(files: list[str]):
     )
 
     assert rewritten_arguments(context) == {"files": {0, 1}}
+
+
+async def test_quarantine_client_uses_its_own_security_config() -> None:
+    """Scoped quarantine calls must never inherit another agent's configured client."""
+    from agent_framework.security import quarantined_llm, set_quarantine_client
+
+    class FakeClient:
+        def __init__(self, name: str) -> None:
+            self.name = name
+            self.calls: list[str] = []
+
+        async def get_response(self, *, messages: Any, **kwargs: Any) -> SimpleNamespace:
+            value = "\n".join(
+                item.text or "" for msg in messages for item in msg.contents if getattr(item, "type", None) == "text"
+            )
+            self.calls.append(value)
+            await asyncio.sleep(0)
+            return SimpleNamespace(text=self.name + " response")
+
+    alpha, beta = FakeClient("A"), FakeClient("B")
+    # The fake exposes only the chat methods exercised by these security tests.
+    a = SecureAgentConfig(quarantine_chat_client=cast(SupportsChatGetResponse, alpha))
+    b = SecureAgentConfig(quarantine_chat_client=cast(SupportsChatGetResponse, beta))
+    empty = SecureAgentConfig(quarantine_chat_client=None)
+
+    async def execute(cfg: SecureAgentConfig, marker: str, *, direct: bool = False) -> None:
+        session = AgentSession(session_id="synthetic-" + marker)
+        ctx = SessionContext(session_id=session.session_id, input_messages=[])
+        if direct:
+            middleware = cfg.get_middleware()
+        else:
+            await cfg.before_run(
+                agent=None,
+                session=session,
+                context=ctx,
+                state=session.state.setdefault(cfg.source_id, {}),
+            )
+            middleware = cast(list[FunctionMiddleware], ctx.middleware[cfg.source_id])
+        tool = next(t for t in cfg.get_tools() if t.name == "quarantined_llm")
+        args = {
+            "prompt": "Summarize",
+            "labelled_data": {
+                "record": {
+                    "content": marker,
+                    "security_label": {"integrity": "untrusted", "confidentiality": "private"},
+                }
+            },
+        }
+        call = FunctionInvocationContext(function=tool, arguments=args, session=session)
+
+        async def invoke(inner: FunctionInvocationContext) -> Any:
+            return await tool.invoke(arguments=inner.arguments, context=inner, skip_parsing=True)
+
+        await FunctionMiddlewarePipeline(*middleware).execute(call, invoke)
+
+    try:
+        await asyncio.gather(execute(a, "ALPHA_MARKER"), execute(b, "BETA_MARKER"))
+        assert any("ALPHA_MARKER" in x for x in alpha.calls)
+        assert not any("ALPHA_MARKER" in x for x in beta.calls)
+        assert any("BETA_MARKER" in x for x in beta.calls)
+        assert not any("BETA_MARKER" in x for x in alpha.calls)
+        await execute(empty, "NO_CLIENT_MARKER")
+        assert not any("NO_CLIENT_MARKER" in x for x in beta.calls)
+        await execute(a, "ALPHA_SECOND_MARKER")
+        assert any("ALPHA_SECOND_MARKER" in x for x in alpha.calls)
+        assert not any("ALPHA_SECOND_MARKER" in x for x in beta.calls)
+
+        # Swapping trackers after config creation is explicitly supported.
+        # Both session-cloned and direct stacks must honor the owning client.
+        a.label_tracker = LabelTrackingFunctionMiddleware()
+        b.label_tracker = LabelTrackingFunctionMiddleware()
+        empty.label_tracker = LabelTrackingFunctionMiddleware()
+        await asyncio.gather(
+            execute(a, "ALPHA_REPLACEMENT_MARKER"),
+            execute(b, "BETA_REPLACEMENT_MARKER"),
+        )
+        assert any("ALPHA_REPLACEMENT_MARKER" in x for x in alpha.calls)
+        assert not any("ALPHA_REPLACEMENT_MARKER" in x for x in beta.calls)
+        assert any("BETA_REPLACEMENT_MARKER" in x for x in beta.calls)
+        assert not any("BETA_REPLACEMENT_MARKER" in x for x in alpha.calls)
+        await execute(empty, "NO_CLIENT_REPLACEMENT_MARKER")
+        assert not any("NO_CLIENT_REPLACEMENT_MARKER" in x for x in beta.calls)
+        await execute(a, "ALPHA_DIRECT_MARKER", direct=True)
+        assert any("ALPHA_DIRECT_MARKER" in x for x in alpha.calls)
+        assert not any("ALPHA_DIRECT_MARKER" in x for x in beta.calls)
+
+        # The provider must also bind trackers supplied by an overriding
+        # get_middleware() implementation that does not call super().
+        class CustomizedConfig(SecureAgentConfig):
+            def get_middleware(self) -> list[FunctionMiddleware]:
+                result: list[FunctionMiddleware] = [self.label_tracker]
+                if self.policy_enforcer is not None:
+                    result.append(self.policy_enforcer)
+                return result
+
+        custom = CustomizedConfig(quarantine_chat_client=cast(SupportsChatGetResponse, alpha))
+        custom.label_tracker = LabelTrackingFunctionMiddleware()
+        await execute(custom, "ALPHA_CUSTOM_STACK_MARKER")
+        assert any("ALPHA_CUSTOM_STACK_MARKER" in x for x in alpha.calls)
+        assert not any("ALPHA_CUSTOM_STACK_MARKER" in x for x in beta.calls)
+
+        # Standalone calls retain the legacy process-global fallback.
+        set_quarantine_client(cast(SupportsChatGetResponse, beta))
+        await quarantined_llm.invoke(
+            arguments={
+                "prompt": "Direct",
+                "labelled_data": {
+                    "record": {
+                        "content": "STANDALONE_MARKER",
+                        "security_label": {"integrity": "untrusted", "confidentiality": "public"},
+                    }
+                },
+            },
+            skip_parsing=True,
+        )
+        assert any("STANDALONE_MARKER" in x for x in beta.calls)
+    finally:
+        set_quarantine_client(None)