diff --git a/python/packages/core/agent_framework/security.py b/python/packages/core/agent_framework/security.py index acc7fbbb81f..0ae5b5e77de 100644 --- a/python/packages/core/agent_framework/security.py +++ b/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 diff --git a/python/packages/core/tests/test_security.py b/python/packages/core/tests/test_security.py index 57aed63a382..46de415b2e9 100644 --- a/python/packages/core/tests/test_security.py +++ b/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)