agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into local-branch-python-add-reset-to-workflow

Browse Source
This commit is contained in:
Tao Chen
2026-06-11 15:28:27 -07:00
Unverified
parent b0d0224ed4 3d5421edc1
commit 96af0cd15a
304 changed files with 12129 additions and 1067 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/CHANGELOG.md
+22 -1
View File
@@ -7,6 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
## [1.8.1] - 2026-06-09
### Added
- **agent-framework-core**: Add MCP client OTel spans per GenAI semantic conventions ([#6349](https://github.com/microsoft/agent-framework/pull/6349))
- **agent-framework-core**: Add MCP long-running task support ([#6319](https://github.com/microsoft/agent-framework/pull/6319))
### Changed
- **agent-framework-claude**: Bump `claude-agent-sdk` to 0.2.87 ([#6248](https://github.com/microsoft/agent-framework/pull/6248))
- **agent-framework-core**: Document checkpoint storage security model and deserialization trust boundaries ([#6295](https://github.com/microsoft/agent-framework/pull/6295))
- **agent-framework-azurefunctions**: Document checkpoint storage security model and deserialization trust boundaries ([#6295](https://github.com/microsoft/agent-framework/pull/6295))
### Fixed
- **agent-framework-core**: Filter MCP tool kwargs to declared params via allowlist ([#6399](https://github.com/microsoft/agent-framework/pull/6399))
- **agent-framework-core**: Fix per-service-call history persistence with server-storing clients ([#6310](https://github.com/microsoft/agent-framework/pull/6310))
- **agent-framework-openai**: Use `getattr` for non-OpenAI provider response compatibility ([#6270](https://github.com/microsoft/agent-framework/pull/6270))
- **agent-framework-foundry-hosting**: Refactor workflow-as-agent pending request handling ([#6259](https://github.com/microsoft/agent-framework/pull/6259))
- **agent-framework-gemini**: Make Gemini honor declarative `outputSchema`, not just JSON mode ([#5893](https://github.com/microsoft/agent-framework/pull/5893))
- **agent-framework-mem0**: Isolate entity retrieval and correct `app_id` payload ([#6242](https://github.com/microsoft/agent-framework/pull/6242))
- **agent-framework-ag-ui**: Match AG-UI approval responses to requested arguments ([#6376](https://github.com/microsoft/agent-framework/pull/6376))
## [1.8.0] - 2026-06-04
### Added
@@ -1169,7 +1189,8 @@ Release candidate for **agent-framework-core** and **agent-framework-azure-ai**
For more information, see the [announcement blog post](https://devblogs.microsoft.com/foundry/introducing-microsoft-agent-framework-the-open-source-engine-for-agentic-ai-apps/).
[Unreleased]: https://github.com/microsoft/agent-framework/compare/python-1.8.0...HEAD
[Unreleased]: https://github.com/microsoft/agent-framework/compare/python-1.8.1...HEAD
[1.8.1]: https://github.com/microsoft/agent-framework/compare/python-1.8.0...python-1.8.1
[1.8.0]: https://github.com/microsoft/agent-framework/compare/python-1.7.0...python-1.8.0
[1.7.0]: https://github.com/microsoft/agent-framework/compare/python-1.6.0...python-1.7.0
[1.6.0]: https://github.com/microsoft/agent-framework/compare/python-1.5.0...python-1.6.0
python/packages/ag-ui/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[project]
name = "agent-framework-ag-ui"
version = "1.0.0rc3"
version = "1.0.0rc4"
description = "AG-UI protocol integration for Agent Framework"
readme = "README.md"
license-files = ["LICENSE"]
@@ -22,7 +22,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.6.0,<2",
"agent-framework-core>=1.8.1,<2",
"ag-ui-protocol>=0.1.16,<0.2",
"fastapi>=0.115.0,<0.133.1",
"uvicorn[standard]>=0.30.0,<1"
python/packages/azurefunctions/agent_framework_azurefunctions/_serialization.py
+18
View File
@@ -14,6 +14,24 @@ This module adds:
- reconstruct_to_type: for HITL responses where external data (without type markers)
needs to be reconstructed to a known type
- resolve_type: resolves 'module:class' type keys to Python types
Security Model
--------------
The underlying Azure Durable Functions storage (Azure Storage account) is the
trusted persistence layer for serialized checkpoint data. The
``RestrictedUnpickler`` in the core encoding module provides defense-in-depth
type filtering, but checkpoint storage itself must be properly access-controlled:
- Ensure the Azure Storage account used by Durable Functions is not publicly
writable and uses appropriate RBAC / shared-access policies.
- Never route untrusted user input directly into ``deserialize_value`` without
first calling :func:`strip_pickle_markers` to neutralize injection of
pickle markers into the data path.
- Configure your checkpoint storage with ``allowed_checkpoint_types`` (or call
``decode_checkpoint_value(..., allowed_types=...)`` directly) to restrict the set of types that can be deserialized.
See :mod:`agent_framework._workflows._checkpoint_encoding` for the full
security model documentation.
"""
from __future__ import annotations
python/packages/azurefunctions/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "Azure Functions integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.0.0b260604"
version = "1.0.0b260609"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -22,7 +22,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.8.0,<2",
"agent-framework-core>=1.8.1,<2",
"agent-framework-durabletask>=1.0.0b260604,<2",
"azure-functions>=1.24.0,<2",
"azure-functions-durable>=1.3.1,<2",
python/packages/claude/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "Claude Agent SDK integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.0.0b260521"
version = "1.0.0b260609"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,7 +23,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.6.0,<2",
"agent-framework-core>=1.8.1,<2",
"claude-agent-sdk>=0.1.36,<0.3",
]
python/packages/core/AGENTS.md
+18
View File
@@ -82,6 +82,7 @@ agent_framework/
- **`MCPStdioTool`** / **`MCPStreamableHTTPTool`** / **`MCPWebsocketTool`** - Transport-specific subclasses.
- **Argument allowlist (`_prepare_call_kwargs`)** - Before each `tools/call`, kwargs are filtered to an **allowlist** built from the tool's declared parameters (`inputSchema.properties`) plus any user-configured extras. Framework runtime kwargs injected through the function-invocation pipeline (e.g. `thread`, `conversation_id`, `chat_options`, `options`, `response_format`) are stripped by default rather than forwarded. A tool that declares no usable `properties` (including schemas with `additionalProperties: true`) forwards only the configured extras. The `_MCP_FRAMEWORK_DENYLIST` is a safety net for framework-named params a server *declares* in its schema (those are dropped); names explicitly opted in via `additional_tool_argument_names` always win. The reserved `_meta` key is extracted as MCP request metadata, never forwarded as an argument.
- **`additional_tool_argument_names`** (constructor arg on all `MCPTool` subclasses) - Opt extra argument names back into the allowlist. Accepts a `Sequence[str]` (applied to every tool) or a `Mapping[str, Sequence[str]]` keyed by **remote tool name**, where the reserved key `"*"` denotes global extras. It is configured only in user code at construction; there is **no per-call/runtime override**, so a model-issued tool call cannot change which names pass through. To use a server that accepts `additionalProperties: true`, list the extra names here and then either (1) manually extend that tool's `inputSchema` (via the `.functions` list after connecting) so the model is prompted to supply them, or (2) supply the values yourself via `function_invocation_kwargs`. If a name is supplied by both the model and `function_invocation_kwargs`, the model-supplied value wins.
- **Sampling guardrails** (`sampling_callback`) - Passing `client=` advertises `SamplingCapability` so the server can send `sampling/createMessage`. Because remote servers are untrusted (confused-deputy risk), the default `sampling_callback` is **deny-by-default** and applies, in order: a per-session rate limit (`sampling_max_requests`, default `_DEFAULT_SAMPLING_MAX_REQUESTS`), an approval gate (`sampling_approval_callback`), and a `maxTokens` cap (`sampling_max_tokens`, default `_DEFAULT_SAMPLING_MAX_TOKENS`). The approval callback (constructor arg on all subclasses; exported type alias `SamplingApprovalCallback`) receives the raw `CreateMessageRequestParams`, may be sync or async, and must return truthy to approve. When it is `None` (the default) every sampling request is denied; pass `lambda params: True` to restore legacy auto-approve as an explicit opt-in. Requests and denials are logged at WARNING (content is not logged). The per-session counter resets in `_reset_session_state`.
- **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields:
- `default_ttl: timedelta | None` — forwarded to the server as `params.task.ttl` (milliseconds). When `None`, the server's default applies.
- `cancel_remote_task_on_local_cancellation: bool = True` — only gates the `CancelledError` path. Abandonment paths (see below) always cancel.
@@ -99,6 +100,23 @@ agent_framework/
- **`FileSearchResult`** / **`FileSearchMatch`** - `SerializationMixin` DTOs returned by `search_files`, carrying the matching file name, a context snippet, and the matching lines with 1-based line numbers.
- **`FileAccessProvider`** - `ContextProvider` that adds shared file-access tools (`file_access_save_file`, `file_access_read_file`, `file_access_delete_file`, `file_access_list_files`, `file_access_search_files`) plus default usage instructions to each invocation. Unlike `MemoryContextProvider`, the store is intentionally shared across sessions and agents.
### Tool Approval Harness (`_harness/_tool_approval.py`)
- **`ToolApprovalMiddleware`** - Experimental opt-in agent middleware that coordinates session-backed approval
rules, heuristic `auto_approval_rules`, queued approval requests, collected approval responses, and
streaming/non-streaming approval prompts. Heuristic callbacks receive the underlying `function_call` content.
- **`ToolApprovalRule`** / **`ToolApprovalState`** - Serializable state models for standing approvals and queued
approval flow. `ToolApprovalRule.arguments is None` means a tool-wide rule; an empty dict `{}` means an exact
no-argument call for `create_always_approve_tool_with_arguments_response`.
- **`create_always_approve_tool_response`** / **`create_always_approve_tool_with_arguments_response`** - Helpers
that return normal `function_approval_response` content with `additional_properties` metadata consumed by
`ToolApprovalMiddleware`. Standing rules for hosted tools include the `server_label` boundary, so same-named tools
on different hosted servers do not share approvals.
- Mixed tool-call batches use a default .NET-style bypass in the function invocation loop: when a session is
available, approval requests for known non-approval-required tools are treated as already approved, hidden, stored
in session state keyed to the visible approval request ids from that batch, and reinjected only when that visible
approval flow resumes.
### Workflows (`_workflows/`)
- **`Workflow`** - Graph-based workflow definition
python/packages/core/agent_framework/__init__.py
+20 -1
View File
@@ -27,6 +27,7 @@ from ._clients import (
SupportsGetEmbeddings,
SupportsImageGenerationTool,
SupportsMCPTool,
SupportsShellTool,
SupportsWebSearchTool,
)
from ._compaction import (
@@ -124,7 +125,16 @@ from ._harness._todo import (
TodoSessionStore,
TodoStore,
)
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPTaskOptions, MCPWebsocketTool
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPTaskOptions, MCPWebsocketTool, SamplingApprovalCallback
from ._harness._tool_approval import (
DEFAULT_TOOL_APPROVAL_SOURCE_ID,
ToolApprovalMiddleware,
ToolApprovalRule,
ToolApprovalRuleCallback,
ToolApprovalState,
create_always_approve_tool_response,
create_always_approve_tool_with_arguments_response,
)
from ._middleware import (
AgentContext,
AgentMiddleware,
@@ -330,6 +340,7 @@ __all__ = [
"DEFAULT_MEMORY_SOURCE_ID",
"DEFAULT_MODE_SOURCE_ID",
"DEFAULT_TODO_SOURCE_ID",
"DEFAULT_TOOL_APPROVAL_SOURCE_ID",
"EXCLUDED_KEY",
"EXCLUDE_REASON_KEY",
"GROUP_ANNOTATION_KEY",
@@ -472,6 +483,7 @@ __all__ = [
"RubricScore",
"RunContext",
"RunnerContext",
"SamplingApprovalCallback",
"SecretString",
"SelectiveToolCallCompactionStrategy",
"SessionContext",
@@ -495,6 +507,7 @@ __all__ = [
"SupportsGetEmbeddings",
"SupportsImageGenerationTool",
"SupportsMCPTool",
"SupportsShellTool",
"SupportsWebSearchTool",
"SwitchCaseEdgeGroup",
"SwitchCaseEdgeGroupCase",
@@ -508,6 +521,10 @@ __all__ = [
"TodoStore",
"TokenBudgetComposedStrategy",
"TokenizerProtocol",
"ToolApprovalMiddleware",
"ToolApprovalRule",
"ToolApprovalRuleCallback",
"ToolApprovalState",
"ToolMode",
"ToolResultCompactionStrategy",
"ToolTypes",
@@ -542,6 +559,8 @@ __all__ = [
"annotate_message_groups",
"apply_compaction",
"chat_middleware",
"create_always_approve_tool_response",
"create_always_approve_tool_with_arguments_response",
"create_edge_runner",
"create_harness_agent",
"detect_media_type_from_base64",
python/packages/core/agent_framework/_clients.py
+30
View File
@@ -819,6 +819,36 @@ class SupportsFileSearchTool(Protocol):
...
@runtime_checkable
class SupportsShellTool(Protocol):
"""Protocol for clients that support shell tools.
This protocol enables runtime checking to determine if a client
supports executing shell commands.
Examples:
.. code-block:: python
from agent_framework import SupportsShellTool
if isinstance(client, SupportsShellTool):
tool = client.get_shell_tool(func=shell.as_function())
agent = ChatAgent(client, tools=[tool])
"""
@staticmethod
def get_shell_tool(**kwargs: Any) -> Any:
"""Create a shell tool configuration.
Keyword Args:
**kwargs: Provider-specific configuration options.
Returns:
A tool configuration ready to pass to ChatAgent.
"""
...
# endregion
python/packages/core/agent_framework/_harness/_agent.py
+131 -25
View File
@@ -15,7 +15,7 @@ from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING, Any
from .._agents import Agent, SupportsAgentRun
from .._clients import SupportsWebSearchTool
from .._clients import SupportsShellTool, SupportsWebSearchTool
from .._compaction import CompactionProvider, ContextWindowCompactionStrategy, ToolResultCompactionStrategy
from .._feature_stage import ExperimentalFeature, experimental
from .._sessions import ContextProvider, HistoryProvider, InMemoryHistoryProvider
@@ -28,6 +28,8 @@ from ._todo import TodoProvider
if TYPE_CHECKING:
from collections.abc import Mapping
from agent_framework_tools.shell import ShellEnvironmentProviderOptions, ShellExecutor
from .._clients import SupportsChatGetResponse
from .._compaction import CompactionStrategy, TokenizerProtocol
from .._middleware import MiddlewareTypes
@@ -66,23 +68,45 @@ def _assemble_instructions(
def _assemble_compaction_provider(
*,
disable_compaction: bool,
max_context_window_tokens: int,
max_output_tokens: int,
max_context_window_tokens: int | None,
max_output_tokens: int | None,
history_source_id: str,
before_compaction_strategy: CompactionStrategy | None,
after_compaction_strategy: CompactionStrategy | None,
tokenizer: TokenizerProtocol | None,
) -> CompactionProvider | None:
"""Build the compaction provider from parameters or defaults."""
"""Build the compaction provider from parameters or defaults.
The token-budget defaults (``ContextWindowCompactionStrategy`` for the before phase and
``ToolResultCompactionStrategy`` for the after phase) are only applied when the token
params are provided. Caller-supplied strategies are always honored. Either phase may end
up ``None``, which ``CompactionProvider`` interprets as "skip that phase".
Returns None when compaction is explicitly disabled, or when neither phase has a strategy
(no custom strategies and no token budget to build the defaults).
"""
if disable_compaction:
return None
before_strategy = before_compaction_strategy or ContextWindowCompactionStrategy(
max_context_window_tokens=max_context_window_tokens,
max_output_tokens=max_output_tokens,
tokenizer=tokenizer,
)
after_strategy = after_compaction_strategy or ToolResultCompactionStrategy(keep_last_tool_call_groups=2)
# Resolve the before-strategy: custom strategy wins; otherwise fall back to the
# token-budget-aware default when token params are available.
before_strategy = before_compaction_strategy
if before_strategy is None and max_context_window_tokens is not None and max_output_tokens is not None:
before_strategy = ContextWindowCompactionStrategy(
max_context_window_tokens=max_context_window_tokens,
max_output_tokens=max_output_tokens,
tokenizer=tokenizer,
)
# Resolve the after-strategy: custom strategy wins; otherwise fall back to the default
# when token params are available.
after_strategy = after_compaction_strategy
if after_strategy is None and max_context_window_tokens is not None and max_output_tokens is not None:
after_strategy = ToolResultCompactionStrategy(keep_last_tool_call_groups=2)
# Nothing to compact in either phase: skip the provider entirely.
if before_strategy is None and after_strategy is None:
return None
return CompactionProvider(
before_strategy=before_strategy,
@@ -106,6 +130,7 @@ def _assemble_context_providers(
skills_paths: Sequence[str] | None,
background_agents: Sequence[SupportsAgentRun] | None,
background_agents_instructions: str | None,
shell_context_provider: ContextProvider | None,
extra_context_providers: Sequence[ContextProvider] | None,
) -> list[ContextProvider]:
"""Assemble the ordered list of context providers."""
@@ -137,6 +162,10 @@ def _assemble_context_providers(
if background_agents:
providers.append(BackgroundAgentsProvider(background_agents, instructions=background_agents_instructions))
# Shell environment provider is opt-in: only added when a shell tool was wired.
if shell_context_provider is not None:
providers.append(shell_context_provider)
# Append any user-supplied additional providers.
if extra_context_providers:
providers.extend(extra_context_providers)
@@ -144,6 +173,50 @@ def _assemble_context_providers(
return providers
def _assemble_shell(
client: SupportsChatGetResponse[Any],
shell_executor: ShellExecutor | None,
shell_environment_provider_options: ShellEnvironmentProviderOptions | None,
) -> tuple[ToolTypes | None, ContextProvider | None]:
"""Build the shell tool and environment provider when a shell executor is supplied.
Returns a ``(tool, provider)`` tuple. Both are ``None`` when no shell executor is
provided, or when the client does not support shell tools (a warning is logged in the
latter case, since the environment provider is not useful without an execution path).
Raises:
TypeError: If ``shell_executor`` does not expose a callable ``as_function()`` method.
"""
if shell_executor is None:
return None, None
# ShellExecutor is a protocol without ``as_function()``, so the
# contract is validated at runtime: a shell tool such as LocalShellTool/DockerShellTool exposes it.
as_function = getattr(shell_executor, "as_function", None)
if not callable(as_function):
raise TypeError(
f"shell_executor must expose a callable 'as_function()' method "
f"(e.g. a LocalShellTool or DockerShellTool from agent-framework-tools), "
f"but got {type(shell_executor).__name__}."
)
if not isinstance(client, SupportsShellTool):
logger.warning(
"Shell tool not available: client %r does not implement SupportsShellTool. "
"Skipping the shell tool and environment provider.",
type(client).__name__,
)
return None, None
# Imported lazily: the shell types live in the separate agent-framework-tools package,
# which depends on core, so core cannot import them at module load time.
from agent_framework_tools.shell import ShellEnvironmentProvider
shell_tool = client.get_shell_tool(func=as_function())
shell_provider = ShellEnvironmentProvider(shell_executor, shell_environment_provider_options)
return shell_tool, shell_provider
HARNESS_AGENT_PROVIDER_NAME = "microsoft.agent_framework.harness"
@@ -157,8 +230,8 @@ def create_harness_agent(
harness_instructions: str | None = None,
agent_instructions: str | None = None,
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
max_context_window_tokens: int,
max_output_tokens: int,
max_context_window_tokens: int | None = None,
max_output_tokens: int | None = None,
history_provider: HistoryProvider | None = None,
disable_compaction: bool = False,
before_compaction_strategy: CompactionStrategy | None = None,
@@ -174,6 +247,8 @@ def create_harness_agent(
skills_paths: Sequence[str] | None = None,
background_agents: Sequence[SupportsAgentRun] | None = None,
background_agents_instructions: str | None = None,
shell_executor: ShellExecutor | None = None,
shell_environment_provider_options: ShellEnvironmentProviderOptions | None = None,
disable_web_search: bool = False,
otel_provider_name: str | None = None,
context_providers: Sequence[ContextProvider] | None = None,
@@ -206,8 +281,6 @@ def create_harness_agent(
agent = create_harness_agent(
OpenAIChatClient(model="gpt-4o"),
max_context_window_tokens=128_000,
max_output_tokens=16_384,
)
session = agent.create_session()
response = await agent.run("Plan a weekend trip to Seattle", session=session)
@@ -243,13 +316,21 @@ def create_harness_agent(
(e.g., "You are a research assistant focused on academic sources.").
tools: Additional tools to include in the agent's toolset.
max_context_window_tokens: Maximum tokens the model's context window supports.
Used to construct the default token-budget-aware compaction strategies. When None
(default) and no custom ``before_compaction_strategy`` / ``after_compaction_strategy``
is provided, compaction is automatically disabled.
max_output_tokens: Maximum output tokens per response.
Used to construct the default compaction strategies and sets a default max_tokens
chat option. When None (default), no default max_tokens option is set, and unless a
custom compaction strategy is provided, compaction is automatically disabled.
history_provider: Custom history provider. When None, an InMemoryHistoryProvider is used.
disable_compaction: When True, skip compaction provider setup.
before_compaction_strategy: Custom before-run compaction strategy.
Defaults to ContextWindowCompactionStrategy (token-budget aware).
after_compaction_strategy: Custom after-run compaction strategy.
Defaults to ToolResultCompactionStrategy.
before_compaction_strategy: Custom before-run compaction strategy. When provided,
compaction runs even if token params are omitted. Defaults to
ContextWindowCompactionStrategy (token-budget aware) when token params are provided.
after_compaction_strategy: Custom after-run compaction strategy. When provided,
compaction runs even if token params are omitted. Defaults to
ToolResultCompactionStrategy when token params are provided.
tokenizer: Custom tokenizer for compaction strategies.
disable_todo: When True, skip the TodoProvider.
todo_provider: Custom TodoProvider instance. Ignored when disable_todo is True.
@@ -270,6 +351,15 @@ def create_harness_agent(
background_agents_instructions: Optional instruction override for the
``BackgroundAgentsProvider``. May include ``{background_agents}`` placeholder
which will be replaced with the agent listing.
shell_executor: Optional shell tool that enables shell command execution. When
provided, the shell tool and a ``ShellEnvironmentProvider`` are automatically
added (provided the client supports shell tools; otherwise a warning is logged
and both are skipped). The object must expose ``as_function()`` and satisfy the
``ShellExecutor`` protocol -- e.g. a ``LocalShellTool`` or ``DockerShellTool`` from
the ``agent-framework-tools`` package. The caller owns the executor's lifecycle.
shell_environment_provider_options: Optional ``ShellEnvironmentProviderOptions``
(from ``agent-framework-tools``) used to customize the ``ShellEnvironmentProvider``
environment probing and instructions. Only used when ``shell_executor`` is provided.
disable_web_search: When True, skip automatic web search tool inclusion.
When False (default), the web search tool is automatically added if the
client implements SupportsWebSearchTool. A warning is logged if the client
@@ -283,14 +373,19 @@ def create_harness_agent(
A fully configured :class:`~agent_framework.Agent` instance.
Raises:
ValueError: If max_context_window_tokens <= 0 or max_output_tokens < 0
or max_output_tokens >= max_context_window_tokens.
ValueError: If max_context_window_tokens is provided and <= 0, or
max_output_tokens is provided and <= 0, or max_output_tokens >=
max_context_window_tokens when both are provided.
"""
if max_context_window_tokens <= 0:
if max_context_window_tokens is not None and max_context_window_tokens <= 0:
raise ValueError("max_context_window_tokens must be positive.")
if max_output_tokens < 0:
raise ValueError("max_output_tokens must be non-negative.")
if max_output_tokens >= max_context_window_tokens:
if max_output_tokens is not None and max_output_tokens <= 0:
raise ValueError("max_output_tokens must be positive.")
if (
max_context_window_tokens is not None
and max_output_tokens is not None
and max_output_tokens >= max_context_window_tokens
):
raise ValueError("max_output_tokens must be less than max_context_window_tokens.")
# Build history provider.
@@ -307,6 +402,13 @@ def create_harness_agent(
tokenizer=tokenizer,
)
# Build the shell tool and environment provider (opt-in via shell_executor).
shell_tool, shell_provider = _assemble_shell(
client,
shell_executor,
shell_environment_provider_options,
)
# Build context providers.
assembled_providers = _assemble_context_providers(
history_provider=resolved_history,
@@ -321,6 +423,7 @@ def create_harness_agent(
skills_paths=skills_paths,
background_agents=background_agents,
background_agents_instructions=background_agents_instructions,
shell_context_provider=shell_provider,
extra_context_providers=context_providers,
)
@@ -338,6 +441,8 @@ def create_harness_agent(
"Set disable_web_search=True to suppress this warning.",
type(client).__name__,
)
if shell_tool is not None:
assembled_tools.append(shell_tool)
if tools is not None:
if isinstance(tools, Sequence):
assembled_tools.extend(tools) # pyright: ignore[reportUnknownArgumentType]
@@ -347,7 +452,8 @@ def create_harness_agent(
# Build default options dict.
default_opts: dict[str, Any] = dict(default_options) if default_options else {}
default_opts.setdefault("max_tokens", max_output_tokens)
if max_output_tokens is not None:
default_opts.setdefault("max_tokens", max_output_tokens)
agent = Agent(
client,
python/packages/core/agent_framework/_harness/_tool_approval.py
+632
View File
@@ -0,0 +1,632 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import copy
import inspect
import json
from asyncio import sleep
from collections.abc import AsyncIterable, Awaitable, Callable, Iterable, Mapping, MutableMapping, Sequence
from typing import Any, Literal, cast
from .._feature_stage import ExperimentalFeature, experimental
from .._middleware import AgentContext, AgentMiddleware
from .._serialization import SerializationMixin
from .._sessions import AgentSession
from .._types import (
AgentResponse,
AgentResponseUpdate,
Content,
FinishReason,
FinishReasonLiteral,
Message,
ResponseStream,
)
DEFAULT_TOOL_APPROVAL_SOURCE_ID = "tool_approval"
_FUNCTION_INVOCATION_BUDGET_STATE_KEY = "_function_invocation_budget_state"
ALWAYS_APPROVE_PROPERTY = "tool_approval"
ALWAYS_APPROVE_SCOPE_PROPERTY = "always_approve"
ALWAYS_APPROVE_TOOL: Literal["tool"] = "tool"
ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS: Literal["tool_with_arguments"] = "tool_with_arguments"
_RULES_KEY = "rules"
_QUEUED_APPROVAL_REQUESTS_KEY = "queued_approval_requests"
_COLLECTED_APPROVAL_RESPONSES_KEY = "collected_approval_responses"
ToolApprovalScope = Literal["tool", "tool_with_arguments"]
ToolApprovalRuleCallback = Callable[[Content], bool | Awaitable[bool]]
def _parse_function_arguments(function_call: Content) -> dict[str, Any]:
arguments = function_call.parse_arguments()
return dict(arguments or {})
def _serialize_argument_value(value: Any) -> str:
return json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
def _serialize_arguments(function_call: Content) -> dict[str, str]:
"""Serialize arguments for exact matching.
``None`` is reserved on :class:`ToolApprovalRule` for tool-wide rules.
An argument-scoped rule for a no-argument call stores ``{}``, so it only
matches future no-argument calls and never becomes a wildcard.
"""
arguments = _parse_function_arguments(function_call)
return {key: _serialize_argument_value(value) for key, value in arguments.items()}
def _server_label(function_call: Content) -> str | None:
"""Return the hosted-tool server boundary for a function call, if present."""
value = function_call.additional_properties.get("server_label")
return value if isinstance(value, str) else None
def _content_from_state(value: Any) -> Content:
if isinstance(value, Content):
return value
if isinstance(value, Mapping):
return Content.from_dict(cast(Mapping[str, Any], value))
raise TypeError(f"Expected Content or mapping state item, got {type(value).__name__}.")
def _contents_from_state(values: Any) -> list[Content]:
if not isinstance(values, list):
return []
state_items = list(cast(Iterable[Any], values))
return [_content_from_state(value) for value in state_items]
def _content_to_state(content: Content) -> dict[str, Any]:
return content.to_dict()
@experimental(feature_id=ExperimentalFeature.HARNESS)
class ToolApprovalRule(SerializationMixin):
"""A standing rule for approving future matching tool calls."""
tool_name: str
arguments: dict[str, str] | None
server_label: str | None
def __init__(
self,
tool_name: str,
arguments: Mapping[str, str] | None = None,
*,
server_label: str | None = None,
) -> None:
"""Initialize a tool approval rule.
Args:
tool_name: The function tool name this rule applies to.
arguments: Optional canonicalized argument values. When omitted, the
rule applies to every call to the tool. Use an empty mapping to
match only no-argument calls.
Keyword Args:
server_label: Optional hosted-tool server boundary. Hosted approvals
only match future approvals from the same server label.
"""
normalized_name = tool_name.strip()
if not normalized_name:
raise ValueError("Tool approval rule tool_name must be a non-empty string.")
self.tool_name = normalized_name
self.arguments = dict(arguments) if arguments is not None else None
self.server_label = server_label
@classmethod
def from_dict(
cls,
value: MutableMapping[str, Any],
/,
*,
dependencies: MutableMapping[str, Any] | None = None,
) -> ToolApprovalRule:
"""Create a rule from serialized state."""
del dependencies
tool_name = value.get("tool_name")
if not isinstance(tool_name, str):
raise ValueError("Tool approval rule tool_name must be a string.")
raw_arguments = value.get("arguments")
if raw_arguments is not None and not isinstance(raw_arguments, Mapping):
raise ValueError("Tool approval rule arguments must be a mapping or None.")
server_label = value.get("server_label")
if server_label is not None and not isinstance(server_label, str):
raise ValueError("Tool approval rule server_label must be a string or None.")
arguments = (
{str(key): str(argument_value) for key, argument_value in cast(Mapping[str, Any], raw_arguments).items()}
if isinstance(raw_arguments, Mapping)
else None
)
return cls(tool_name=tool_name, arguments=arguments, server_label=server_label)
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
"""Serialize the rule."""
exclude = exclude or set()
payload: dict[str, Any] = {"tool_name": self.tool_name}
if "type" not in exclude:
payload["type"] = self._get_type_identifier()
if self.arguments is not None or not exclude_none:
payload["arguments"] = self.arguments
if self.server_label is not None or not exclude_none:
payload["server_label"] = self.server_label
return payload
@experimental(feature_id=ExperimentalFeature.HARNESS)
class ToolApprovalState(SerializationMixin):
"""Session-backed state used by :class:`ToolApprovalMiddleware`."""
rules: list[ToolApprovalRule]
queued_approval_requests: list[Content]
collected_approval_responses: list[Content]
def __init__(
self,
*,
rules: Sequence[ToolApprovalRule | Mapping[str, Any]] | None = None,
queued_approval_requests: Sequence[Content | Mapping[str, Any]] | None = None,
collected_approval_responses: Sequence[Content | Mapping[str, Any]] | None = None,
) -> None:
"""Initialize approval state."""
self.rules = [
rule if isinstance(rule, ToolApprovalRule) else ToolApprovalRule.from_dict(dict(rule))
for rule in (rules or [])
]
self.queued_approval_requests = [
item if isinstance(item, Content) else Content.from_dict(item) for item in (queued_approval_requests or [])
]
self.collected_approval_responses = [
item if isinstance(item, Content) else Content.from_dict(item)
for item in (collected_approval_responses or [])
]
@classmethod
def from_dict(
cls,
value: MutableMapping[str, Any],
/,
*,
dependencies: MutableMapping[str, Any] | None = None,
) -> ToolApprovalState:
"""Create state from serialized state."""
del dependencies
return cls(
rules=cast(Sequence[Mapping[str, Any]], value.get("rules", [])),
queued_approval_requests=cast(Sequence[Mapping[str, Any]], value.get("queued_approval_requests", [])),
collected_approval_responses=cast(
Sequence[Mapping[str, Any]],
value.get("collected_approval_responses", []),
),
)
def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]:
"""Serialize state."""
del exclude_none
exclude = exclude or set()
payload: dict[str, Any] = {
"rules": [rule.to_dict() for rule in self.rules],
"queued_approval_requests": [_content_to_state(item) for item in self.queued_approval_requests],
"collected_approval_responses": [_content_to_state(item) for item in self.collected_approval_responses],
}
if "type" not in exclude:
payload["type"] = self._get_type_identifier()
return payload
def create_always_approve_tool_response(request: Content, *, reason: str | None = None) -> Content:
"""Create an approval response that records a standing rule for the whole tool.
Args:
request: The ``function_approval_request`` content to approve.
Keyword Args:
reason: Optional approval reason stored in ``additional_properties``.
Returns:
A ``function_approval_response`` with metadata consumed by
:class:`ToolApprovalMiddleware`.
"""
return _create_always_approve_response(request, ALWAYS_APPROVE_TOOL, reason=reason)
def create_always_approve_tool_with_arguments_response(request: Content, *, reason: str | None = None) -> Content:
"""Create an approval response that records a standing rule for the tool and exact arguments."""
return _create_always_approve_response(request, ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS, reason=reason)
def _create_always_approve_response(request: Content, scope: ToolApprovalScope, *, reason: str | None) -> Content:
response = request.to_function_approval_response(approved=True)
metadata: dict[str, Any] = {ALWAYS_APPROVE_SCOPE_PROPERTY: scope}
if reason is not None:
metadata["reason"] = reason
response.additional_properties[ALWAYS_APPROVE_PROPERTY] = metadata
return response
def _get_state(session: AgentSession, *, source_id: str) -> ToolApprovalState:
raw_state = session.state.get(source_id)
if isinstance(raw_state, ToolApprovalState):
return raw_state
if isinstance(raw_state, MutableMapping):
raw_state_mapping = cast(MutableMapping[str, Any], raw_state)
return ToolApprovalState(
rules=cast(Sequence[Mapping[str, Any]], raw_state_mapping.get(_RULES_KEY, [])),
queued_approval_requests=_contents_from_state(raw_state_mapping.get(_QUEUED_APPROVAL_REQUESTS_KEY, [])),
collected_approval_responses=_contents_from_state(
raw_state_mapping.get(_COLLECTED_APPROVAL_RESPONSES_KEY, []),
),
)
if raw_state is not None:
raise TypeError(f"Session state for {source_id!r} must be a mapping, got {type(raw_state).__name__}.")
state = ToolApprovalState()
session.state[source_id] = state.to_dict(exclude={"type"})
return state
def _save_state(session: AgentSession, state: ToolApprovalState, *, source_id: str) -> None:
serialized = state.to_dict(exclude={"type"})
existing = session.state.get(source_id)
if isinstance(existing, MutableMapping):
for key, value in cast(MutableMapping[str, Any], existing).items():
if key not in serialized and key != "type":
serialized[key] = value
session.state[source_id] = serialized
def _rule_exists(rules: Sequence[ToolApprovalRule], new_rule: ToolApprovalRule) -> bool:
for rule in rules:
if rule.tool_name != new_rule.tool_name:
continue
if rule.server_label != new_rule.server_label:
continue
if rule.arguments == new_rule.arguments:
return True
return False
def _add_rule_if_missing(state: ToolApprovalState, rule: ToolApprovalRule) -> None:
if not _rule_exists(state.rules, rule):
state.rules.append(rule)
def _function_call_from_request(request: Content) -> Content | None:
function_call = request.function_call
if function_call is None or function_call.type != "function_call" or function_call.name is None:
return None
return function_call
def _arguments_match(rule_arguments: Mapping[str, str], function_call: Content) -> bool:
call_arguments = _serialize_arguments(function_call) or {}
if len(rule_arguments) != len(call_arguments):
return False
return all(call_arguments.get(key) == value for key, value in rule_arguments.items())
def _matches_rule(request: Content, rules: Sequence[ToolApprovalRule]) -> bool:
function_call = _function_call_from_request(request)
if function_call is None:
return False
for rule in rules:
if rule.tool_name != function_call.name:
continue
if rule.server_label != _server_label(function_call):
continue
if rule.arguments is None:
return True
if _arguments_match(rule.arguments, function_call):
return True
return False
def _get_always_approve_scope(response: Content) -> ToolApprovalScope | None:
metadata = response.additional_properties.get(ALWAYS_APPROVE_PROPERTY)
if not isinstance(metadata, Mapping):
return None
metadata_mapping = cast(Mapping[str, Any], metadata)
scope = metadata_mapping.get(ALWAYS_APPROVE_SCOPE_PROPERTY)
if scope == ALWAYS_APPROVE_TOOL:
return ALWAYS_APPROVE_TOOL
if scope == ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS:
return ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS
return None
def _clone_without_always_approve_metadata(response: Content) -> Content:
cloned = copy.deepcopy(response)
cloned.additional_properties.pop(ALWAYS_APPROVE_PROPERTY, None)
return cloned
@experimental(feature_id=ExperimentalFeature.HARNESS)
class ToolApprovalMiddleware(AgentMiddleware):
"""Coordinate standing tool approvals and queued approval prompts for an agent.
This middleware is opt-in and requires callers to run the agent with an
:class:`AgentSession`, because approval rules and queued requests are stored
in session state.
"""
def __init__(
self,
*,
source_id: str = DEFAULT_TOOL_APPROVAL_SOURCE_ID,
auto_approval_rules: Sequence[ToolApprovalRuleCallback] | None = None,
) -> None:
"""Initialize the middleware.
Keyword Args:
source_id: Session-state key used by this middleware.
auto_approval_rules: Optional callbacks that can auto-approve a
``function_call``. Each callback receives the function-call
content and returns ``True`` to approve it.
"""
self.source_id = source_id
self.auto_approval_rules = tuple(auto_approval_rules or ())
async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None:
"""Process one agent invocation."""
if context.session is None:
raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.")
state = _get_state(context.session, source_id=self.source_id)
context.client_kwargs.setdefault(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, {})
context.messages = self._prepare_inbound_messages(context.messages, state)
await self._drain_auto_approvable_queue(state)
if next_queued := self._pop_next_queued_request(state):
_save_state(context.session, state, source_id=self.source_id)
context.result = self._response_for_queued_request(next_queued, stream=context.stream)
return
if context.stream:
context.result = self._process_stream(context, call_next, state)
return
while True:
context.messages = self._inject_collected_responses(context.messages, state)
state_changed = bool(state.collected_approval_responses)
state.collected_approval_responses.clear()
if state_changed:
_save_state(context.session, state, source_id=self.source_id)
await call_next()
if isinstance(context.result, ResponseStream):
return
if context.result is None:
_save_state(context.session, state, source_id=self.source_id)
return
all_auto_approved = await self._process_outbound_messages(context.result.messages, state)
_save_state(context.session, state, source_id=self.source_id)
if not all_auto_approved:
return
context.messages = []
context.result = None
def _response_for_queued_request(
self,
request: Content,
*,
stream: bool,
) -> AgentResponse | ResponseStream[AgentResponseUpdate, AgentResponse]:
if not stream:
return AgentResponse(messages=[Message(role="assistant", contents=[request])])
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
await sleep(0)
yield AgentResponseUpdate(role="assistant", contents=[request])
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
def _process_stream(
self,
context: AgentContext,
call_next: Callable[[], Awaitable[None]],
state: ToolApprovalState,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
if context.session is None:
raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.")
while True:
context.messages = self._inject_collected_responses(context.messages, state)
state_changed = bool(state.collected_approval_responses)
state.collected_approval_responses.clear()
if state_changed:
_save_state(context.session, state, source_id=self.source_id)
await call_next()
if not isinstance(context.result, ResponseStream):
raise ValueError("Streaming ToolApprovalMiddleware requires a ResponseStream result.")
approval_requests: list[Content] = []
async for update in context.result:
approval_contents = [
content for content in update.contents if content.type == "function_approval_request"
]
if not approval_contents:
yield update
continue
approval_requests.extend(approval_contents)
remaining_contents = [
content for content in update.contents if content.type != "function_approval_request"
]
if remaining_contents:
raw_finish_reason = update.finish_reason
finish_reason: FinishReasonLiteral | FinishReason | None
if isinstance(raw_finish_reason, str):
finish_reason = FinishReason(raw_finish_reason)
else:
finish_reason = cast(FinishReasonLiteral | FinishReason | None, raw_finish_reason)
yield AgentResponseUpdate(
contents=remaining_contents,
role=update.role,
author_name=update.author_name,
agent_id=update.agent_id,
response_id=update.response_id,
message_id=update.message_id,
created_at=update.created_at,
finish_reason=finish_reason,
continuation_token=update.continuation_token,
additional_properties=update.additional_properties,
raw_representation=update.raw_representation,
)
await context.result.get_final_response()
if not approval_requests:
return
response_messages = [Message(role="assistant", contents=approval_requests)]
all_auto_approved = await self._process_outbound_messages(response_messages, state)
_save_state(context.session, state, source_id=self.source_id)
if not all_auto_approved:
for message in response_messages:
if message.contents:
yield AgentResponseUpdate(role=message.role, contents=message.contents)
return
context.messages = []
context.result = None
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
def _prepare_inbound_messages(self, messages: Sequence[Message], state: ToolApprovalState) -> list[Message]:
prepared: list[Message] = []
for message in messages:
replacement_contents: list[Content] = []
changed = False
for content in message.contents:
if content.type == "function_approval_response":
replacement = self._handle_inbound_approval_response(content, state)
state.collected_approval_responses.append(replacement)
changed = True
continue
replacement_contents.append(content)
if not changed:
prepared.append(message)
continue
if replacement_contents:
cloned = copy.copy(message)
cloned.contents = replacement_contents
prepared.append(cloned)
return prepared
def _handle_inbound_approval_response(self, response: Content, state: ToolApprovalState) -> Content:
scope = _get_always_approve_scope(response)
if scope is None or not response.approved:
return response
function_call = response.function_call
if function_call is not None and function_call.type == "function_call" and function_call.name is not None:
if scope == ALWAYS_APPROVE_TOOL:
_add_rule_if_missing(
state,
ToolApprovalRule(
tool_name=function_call.name,
server_label=_server_label(function_call),
),
)
else:
_add_rule_if_missing(
state,
ToolApprovalRule(
tool_name=function_call.name,
arguments=_serialize_arguments(function_call),
server_label=_server_label(function_call),
),
)
return _clone_without_always_approve_metadata(response)
def _inject_collected_responses(self, messages: Sequence[Message], state: ToolApprovalState) -> list[Message]:
if not state.collected_approval_responses:
return list(messages)
return [Message(role="user", contents=list(state.collected_approval_responses)), *messages]
async def _drain_auto_approvable_queue(self, state: ToolApprovalState) -> None:
remaining: list[Content] = []
for request in state.queued_approval_requests:
if _matches_rule(request, state.rules) or await self._matches_auto_rule(request):
state.collected_approval_responses.append(request.to_function_approval_response(approved=True))
continue
remaining.append(request)
state.queued_approval_requests = remaining
def _pop_next_queued_request(self, state: ToolApprovalState) -> Content | None:
if not state.queued_approval_requests:
return None
return state.queued_approval_requests.pop(0)
async def _process_outbound_messages(self, messages: list[Message], state: ToolApprovalState) -> bool:
approval_requests = [
content
for message in messages
for content in message.contents
if content.type == "function_approval_request"
]
if not approval_requests:
return False
auto_approved: set[int] = set()
unresolved: list[Content] = []
for request in approval_requests:
if _matches_rule(request, state.rules) or await self._matches_auto_rule(request):
state.collected_approval_responses.append(request.to_function_approval_response(approved=True))
auto_approved.add(id(request))
else:
unresolved.append(request)
if not auto_approved and len(unresolved) <= 1:
return False
queued_ids: set[int] = set()
for request in unresolved[1:]:
queued_ids.add(id(request))
state.queued_approval_requests.append(request)
remove_ids = auto_approved | queued_ids
self._remove_approval_requests(messages, remove_ids)
return not unresolved
@staticmethod
def _remove_approval_requests(messages: list[Message], remove_ids: set[int]) -> None:
for message_index in range(len(messages) - 1, -1, -1):
message = messages[message_index]
filtered = [
content
for content in message.contents
if content.type != "function_approval_request" or id(content) not in remove_ids
]
if len(filtered) == len(message.contents):
continue
if filtered:
message.contents = filtered
else:
messages.pop(message_index)
async def _matches_auto_rule(self, request: Content) -> bool:
function_call = _function_call_from_request(request)
if function_call is None:
return False
for rule in self.auto_approval_rules:
result = rule(function_call)
if inspect.isawaitable(result):
result = await result
if result:
return True
return False
__all__ = [
"ALWAYS_APPROVE_PROPERTY",
"ALWAYS_APPROVE_SCOPE_PROPERTY",
"ALWAYS_APPROVE_TOOL",
"ALWAYS_APPROVE_TOOL_WITH_ARGUMENTS",
"DEFAULT_TOOL_APPROVAL_SOURCE_ID",
"ToolApprovalMiddleware",
"ToolApprovalRule",
"ToolApprovalRuleCallback",
"ToolApprovalState",
"create_always_approve_tool_response",
"create_always_approve_tool_with_arguments_response",
]
python/packages/core/agent_framework/_mcp.py
+213 -13
View File
@@ -16,6 +16,7 @@ from contextlib import AsyncExitStack, _AsyncGeneratorContextManager # type: ig
from dataclasses import dataclass
from datetime import timedelta
from functools import partial
from inspect import isawaitable
from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast
from opentelemetry import propagate
@@ -99,6 +100,22 @@ _mcp_call_headers: contextvars.ContextVar[dict[str, str]] = contextvars.ContextV
MCP_DEFAULT_TIMEOUT = 30
MCP_DEFAULT_SSE_READ_TIMEOUT = 60 * 5
# Default safety limits applied to server-initiated MCP sampling requests
# (``sampling/createMessage``). MCP servers are untrusted third parties, so the
# default ``sampling_callback`` denies requests unless an approval callback is
# supplied, and bounds the cost of any approved request.
# - ``_DEFAULT_SAMPLING_MAX_TOKENS`` clamps the server-requested ``maxTokens``.
# - ``_DEFAULT_SAMPLING_MAX_REQUESTS`` caps the number of sampling requests per
# session connection (the counter resets on reconnect).
_DEFAULT_SAMPLING_MAX_TOKENS = 4096
_DEFAULT_SAMPLING_MAX_REQUESTS = 25
# A user-supplied gate invoked before each server-initiated sampling request is
# forwarded to the chat client. It receives the raw ``CreateMessageRequestParams``
# and returns (or awaits to) a truthy value to approve the request or a falsy
# value to deny it. Both synchronous and asynchronous callables are supported.
SamplingApprovalCallback = Callable[["types.CreateMessageRequestParams"], "bool | Coroutine[Any, Any, bool]"]
# region: Helpers
LOG_LEVEL_MAPPING: dict[str, int] = {
@@ -345,6 +362,9 @@ class MCPTool:
session: ClientSession | None = None,
request_timeout: int | None = None,
client: SupportsChatGetResponse | None = None,
sampling_approval_callback: SamplingApprovalCallback | None = None,
sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS,
sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS,
additional_properties: dict[str, Any] | None = None,
task_options: MCPTaskOptions | None = None,
additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
@@ -359,7 +379,13 @@ class MCPTool:
name: The name of the MCP tool.
description: A description of the MCP tool.
approval_mode: Whether approval is required to run tools.
allowed_tools: A collection of tool names to allow.
allowed_tools: Optional allow-list of MCP tool names to expose as functions.
``None`` (the default) exposes every tool advertised by the MCP server.
A non-empty collection exposes only the tools whose names appear in it.
An empty collection (``[]``) exposes no tools — if you simply want to
disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is
useful as a runtime guard or when you want to load tool metadata for
inspection without exposing the tools for invocation.
tool_name_prefix: Optional prefix to prepend to exposed MCP function names.
load_tools: Whether to load tools from the MCP server.
parse_tool_results: An optional callable with signature
@@ -378,6 +404,20 @@ class MCPTool:
session: An existing MCP client session to use.
request_timeout: Timeout in seconds for MCP requests.
client: A chat client for sampling callbacks.
sampling_approval_callback: Optional gate invoked before each server-initiated
``sampling/createMessage`` request is forwarded to ``client``. It receives the
raw ``CreateMessageRequestParams`` and may be synchronous or asynchronous;
returning a truthy value approves the request and a falsy value denies it. When
``None`` (the default), every sampling request is **denied** because MCP servers
are untrusted third parties (confused-deputy risk). To restore the legacy
auto-approve behavior, pass ``lambda params: True`` as an explicit, conscious
opt-in.
sampling_max_tokens: Upper bound applied to the server-requested ``maxTokens`` for an
approved sampling request. The effective value is ``min(requested, cap)``. Set to
``None`` to disable the cap. Defaults to ``_DEFAULT_SAMPLING_MAX_TOKENS``.
sampling_max_requests: Maximum number of sampling requests allowed per session
connection; further requests are rejected. The counter resets on reconnect. Set
to ``None`` to disable the limit. Defaults to ``_DEFAULT_SAMPLING_MAX_REQUESTS``.
additional_properties: Additional properties for the tool.
task_options: Options controlling how long-running MCP tasks are driven for
tools that advertise ``execution.taskSupport == "required"``. When ``None``,
@@ -410,6 +450,10 @@ class MCPTool:
self.session = session
self.request_timeout = request_timeout
self.client = client
self.sampling_approval_callback = sampling_approval_callback
self.sampling_max_tokens = sampling_max_tokens
self.sampling_max_requests = sampling_max_requests
self._sampling_request_count = 0
self._functions: list[FunctionTool] = []
self._tool_call_meta_by_name: dict[str, dict[str, Any]] = {}
self._tool_task_support_by_name: dict[str, str] = {}
@@ -539,6 +583,9 @@ class MCPTool:
case _:
result.append(Content.from_text(str(item)))
if mcp_type.structuredContent is not None:
result.append(Content.from_text(json.dumps(mcp_type.structuredContent, default=str)))
if not result:
result.append(Content.from_text("null"))
return result
@@ -698,7 +745,7 @@ class MCPTool:
@property
def functions(self) -> list[FunctionTool]:
"""Get the list of functions that are allowed."""
if not self.allowed_tools:
if self.allowed_tools is None:
return self._functions
allowed_names = set(self.allowed_tools)
filtered_functions: list[FunctionTool] = []
@@ -840,6 +887,7 @@ class MCPTool:
self._supports_prompts = True
self._supports_logging = None
self._ping_available = True
self._sampling_request_count = 0
def _set_server_capabilities(self, capabilities: types.ServerCapabilities | None) -> None:
self._server_capabilities = capabilities
@@ -994,6 +1042,49 @@ class MCPTool:
except Exception as exc:
logger.warning("Failed to set log level to %s", logger.level, exc_info=exc)
async def _sampling_request_approved(self, params: types.CreateMessageRequestParams) -> bool:
"""Run the configured sampling approval gate.
Returns ``True`` only when an approval callback is configured and approves the request.
When no callback is set, the request is denied (safe default for untrusted servers).
"""
callback = self.sampling_approval_callback
if callback is None:
logger.warning(
"Denying MCP sampling request from '%s': no 'sampling_approval_callback' configured.",
self.name,
)
return False
try:
outcome = callback(params)
if isawaitable(outcome):
outcome = await outcome
except Exception as ex:
logger.warning(
"Denying MCP sampling request from '%s': approval callback raised %s.",
self.name,
ex,
exc_info=True,
)
return False
approved = bool(outcome)
if not approved:
logger.warning("MCP sampling request from '%s' was denied by the approval callback.", self.name)
return approved
def _capped_sampling_max_tokens(self, requested: int) -> int:
"""Clamp the server-requested ``maxTokens`` to ``sampling_max_tokens`` when configured."""
cap = self.sampling_max_tokens
if cap is not None and requested > cap:
logger.warning(
"Capping MCP sampling maxTokens for '%s' from %d to %d.",
self.name,
requested,
cap,
)
return cap
return requested
async def sampling_callback(
self,
context: RequestContext[ClientSession, Any],
@@ -1001,20 +1092,32 @@ class MCPTool:
) -> types.CreateMessageResult | types.ErrorData:
"""Callback function for sampling.
This function is called when the MCP server needs to get a message completed.
It uses the configured chat client to generate responses.
This function is called when the MCP server sends a ``sampling/createMessage``
request. It enforces safety guardrails and, if the request is approved, uses the
configured chat client to generate a response.
Safety:
MCP servers are untrusted third parties, so forwarding server-controlled prompts
to the chat client without review is a confused-deputy risk. This callback
therefore applies, in order: a per-session rate limit
(``sampling_max_requests``), an approval gate (``sampling_approval_callback``,
which **denies by default** when not configured), and a ``maxTokens`` cap
(``sampling_max_tokens``). To allow sampling, pass a ``sampling_approval_callback``
that returns a truthy value (use ``lambda params: True`` to auto-approve as an
explicit opt-in).
Note:
This is a simple version of this function. It can be overridden to allow
more complex sampling. It gets added to the session at initialization time,
so overriding it is the best way to customize this behavior.
This is the default implementation. It can be overridden to allow more complex
sampling. It gets added to the session at initialization time, so overriding it is
the best way to customize this behavior.
Args:
context: The request context from the MCP server.
params: The message creation request parameters.
Returns:
Either a CreateMessageResult with the generated message or ErrorData if generation fails.
Either a CreateMessageResult with the generated message or ErrorData if the request
is denied, rate limited, or generation fails.
"""
from mcp import types
@@ -1023,7 +1126,38 @@ class MCPTool:
code=types.INTERNAL_ERROR,
message="No chat client available. Please set a chat client.",
)
logger.debug("Sampling callback called with params: %s", params)
logger.warning(
"MCP server '%s' sent a sampling/createMessage request (%d message(s), maxTokens=%s).",
self.name,
len(params.messages),
params.maxTokens,
)
if self.sampling_max_requests is not None:
if self._sampling_request_count >= self.sampling_max_requests:
logger.warning(
"Denying MCP sampling request from '%s': per-session limit of %d reached.",
self.name,
self.sampling_max_requests,
)
return types.ErrorData(
code=types.INVALID_REQUEST,
message="Sampling rate limit exceeded for this MCP session.",
)
self._sampling_request_count += 1
if not await self._sampling_request_approved(params):
if self.sampling_approval_callback is None:
message = (
"Sampling request denied. MCP sampling is disabled by default for untrusted "
"servers; provide a 'sampling_approval_callback' that approves the request to "
"enable it."
)
else:
message = "Sampling request denied by the 'sampling_approval_callback'."
return types.ErrorData(code=types.INVALID_REQUEST, message=message)
messages: list[Message] = []
for msg in params.messages:
messages.append(self._parse_message_from_mcp(msg))
@@ -1045,7 +1179,7 @@ class MCPTool:
if params.temperature is not None:
options["temperature"] = params.temperature
options["max_tokens"] = params.maxTokens
options["max_tokens"] = self._capped_sampling_max_tokens(params.maxTokens)
if params.stopSequences is not None:
options["stop"] = params.stopSequences
@@ -2219,6 +2353,9 @@ class MCPStdioTool(MCPTool):
env: dict[str, str] | None = None,
encoding: str | None = None,
client: SupportsChatGetResponse | None = None,
sampling_approval_callback: SamplingApprovalCallback | None = None,
sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS,
sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS,
additional_properties: dict[str, Any] | None = None,
task_options: MCPTaskOptions | None = None,
additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
@@ -2260,12 +2397,28 @@ class MCPStdioTool(MCPTool):
- A dict with keys `always_require_approval` or `never_require_approval`,
followed by a sequence of strings with the names of the relevant tools.
A tool should not be listed in both, if so, it will require approval.
allowed_tools: A list of tools that are allowed to use this tool.
allowed_tools: Optional allow-list of MCP tool names to expose as functions.
``None`` (the default) exposes every tool advertised by the MCP server.
A non-empty collection exposes only the tools whose names appear in it.
An empty collection (``[]``) exposes no tools — if you simply want to
disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is
useful as a runtime guard or when you want to load tool metadata for
inspection without exposing the tools for invocation.
additional_properties: Additional properties.
args: The arguments to pass to the command.
env: The environment variables to set for the command.
encoding: The encoding to use for the command output.
client: The chat client to use for sampling.
sampling_approval_callback: Optional gate run before each server-initiated
``sampling/createMessage`` request reaches ``client``. Receives the raw
``CreateMessageRequestParams`` (sync or async); a truthy return approves the
request, a falsy return denies it. When ``None`` (the default) every sampling
request is **denied**, since MCP servers are untrusted (confused-deputy risk).
Pass ``lambda params: True`` to auto-approve as an explicit opt-in.
sampling_max_tokens: Cap applied to an approved request's ``maxTokens``
(``min(requested, cap)``); ``None`` disables it.
sampling_max_requests: Per-session cap on the number of sampling requests; further
requests are rejected. Resets on reconnect. ``None`` disables it.
task_options: Options for tools that advertise
``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`.
additional_tool_argument_names: Extra argument names to forward to the MCP server in
@@ -2300,6 +2453,9 @@ class MCPStdioTool(MCPTool):
request_timeout=request_timeout,
task_options=task_options,
additional_tool_argument_names=additional_tool_argument_names,
sampling_approval_callback=sampling_approval_callback,
sampling_max_tokens=sampling_max_tokens,
sampling_max_requests=sampling_max_requests,
)
self.command = command
self.args = args or []
@@ -2375,6 +2531,9 @@ class MCPStreamableHTTPTool(MCPTool):
allowed_tools: Collection[str] | None = None,
terminate_on_close: bool | None = None,
client: SupportsChatGetResponse | None = None,
sampling_approval_callback: SamplingApprovalCallback | None = None,
sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS,
sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS,
additional_properties: dict[str, Any] | None = None,
http_client: AsyncClient | None = None,
header_provider: Callable[[dict[str, Any]], dict[str, str]] | None = None,
@@ -2419,10 +2578,26 @@ class MCPStreamableHTTPTool(MCPTool):
- A dict with keys `always_require_approval` or `never_require_approval`,
followed by a sequence of strings with the names of the relevant tools.
A tool should not be listed in both, if so, it will require approval.
allowed_tools: A list of tools that are allowed to use this tool.
allowed_tools: Optional allow-list of MCP tool names to expose as functions.
``None`` (the default) exposes every tool advertised by the MCP server.
A non-empty collection exposes only the tools whose names appear in it.
An empty collection (``[]``) exposes no tools — if you simply want to
disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is
useful as a runtime guard or when you want to load tool metadata for
inspection without exposing the tools for invocation.
additional_properties: Additional properties.
terminate_on_close: Close the transport when the MCP client is terminated.
client: The chat client to use for sampling.
sampling_approval_callback: Optional gate run before each server-initiated
``sampling/createMessage`` request reaches ``client``. Receives the raw
``CreateMessageRequestParams`` (sync or async); a truthy return approves the
request, a falsy return denies it. When ``None`` (the default) every sampling
request is **denied**, since MCP servers are untrusted (confused-deputy risk).
Pass ``lambda params: True`` to auto-approve as an explicit opt-in.
sampling_max_tokens: Cap applied to an approved request's ``maxTokens``
(``min(requested, cap)``); ``None`` disables it.
sampling_max_requests: Per-session cap on the number of sampling requests; further
requests are rejected. Resets on reconnect. ``None`` disables it.
http_client: Optional asyncClient to use. If not provided, the
``streamable_http_client`` API will create and manage a default client.
To configure headers, timeouts, or other HTTP client settings, create
@@ -2466,6 +2641,9 @@ class MCPStreamableHTTPTool(MCPTool):
request_timeout=request_timeout,
task_options=task_options,
additional_tool_argument_names=additional_tool_argument_names,
sampling_approval_callback=sampling_approval_callback,
sampling_max_tokens=sampling_max_tokens,
sampling_max_requests=sampling_max_requests,
)
self.url = url
self.terminate_on_close = terminate_on_close
@@ -2590,6 +2768,9 @@ class MCPWebsocketTool(MCPTool):
approval_mode: (Literal["always_require", "never_require"] | MCPSpecificApproval | None) = None,
allowed_tools: Collection[str] | None = None,
client: SupportsChatGetResponse | None = None,
sampling_approval_callback: SamplingApprovalCallback | None = None,
sampling_max_tokens: int | None = _DEFAULT_SAMPLING_MAX_TOKENS,
sampling_max_requests: int | None = _DEFAULT_SAMPLING_MAX_REQUESTS,
additional_properties: dict[str, Any] | None = None,
task_options: MCPTaskOptions | None = None,
additional_tool_argument_names: Sequence[str] | Mapping[str, Sequence[str]] | None = None,
@@ -2632,9 +2813,25 @@ class MCPWebsocketTool(MCPTool):
- A dict with keys `always_require_approval` or `never_require_approval`,
followed by a sequence of strings with the names of the relevant tools.
A tool should not be listed in both, if so, it will require approval.
allowed_tools: A list of tools that are allowed to use this tool.
allowed_tools: Optional allow-list of MCP tool names to expose as functions.
``None`` (the default) exposes every tool advertised by the MCP server.
A non-empty collection exposes only the tools whose names appear in it.
An empty collection (``[]``) exposes no tools — if you simply want to
disable tool execution, prefer ``load_tools=False`` instead. ``[]`` is
useful as a runtime guard or when you want to load tool metadata for
inspection without exposing the tools for invocation.
additional_properties: Additional properties.
client: The chat client to use for sampling.
sampling_approval_callback: Optional gate run before each server-initiated
``sampling/createMessage`` request reaches ``client``. Receives the raw
``CreateMessageRequestParams`` (sync or async); a truthy return approves the
request, a falsy return denies it. When ``None`` (the default) every sampling
request is **denied**, since MCP servers are untrusted (confused-deputy risk).
Pass ``lambda params: True`` to auto-approve as an explicit opt-in.
sampling_max_tokens: Cap applied to an approved request's ``maxTokens``
(``min(requested, cap)``); ``None`` disables it.
sampling_max_requests: Per-session cap on the number of sampling requests; further
requests are rejected. Resets on reconnect. ``None`` disables it.
task_options: Options for tools that advertise
``execution.taskSupport == "required"``. See :class:`MCPTaskOptions`.
additional_tool_argument_names: Extra argument names to forward to the MCP server in
@@ -2669,6 +2866,9 @@ class MCPWebsocketTool(MCPTool):
request_timeout=request_timeout,
task_options=task_options,
additional_tool_argument_names=additional_tool_argument_names,
sampling_approval_callback=sampling_approval_callback,
sampling_max_tokens=sampling_max_tokens,
sampling_max_requests=sampling_max_requests,
)
self.url = url
self._client_kwargs = kwargs
python/packages/core/agent_framework/_tools.py
+190 -15
View File
@@ -90,6 +90,9 @@ logger = logging.getLogger("agent_framework")
DEFAULT_MAX_ITERATIONS: Final[int] = 40
DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3
SHELL_TOOL_KIND_VALUE: Final[str] = "shell"
_TOOL_APPROVAL_STATE_KEY: Final[str] = "tool_approval"
_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY: Final[str] = "already_approved_approval_request_groups"
_FUNCTION_INVOCATION_BUDGET_STATE_KEY: Final[str] = "_function_invocation_budget_state"
ApprovalMode: TypeAlias = Literal["always_require", "never_require"]
ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]")
ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel)
@@ -1685,15 +1688,15 @@ async def _try_execute_function_calls(
# The live tools list (when tools is the run-local list) is exposed on the
# FunctionInvocationContext so tools can add/remove tools during the run.
live_tools: list[ToolTypes] | None = cast("list[ToolTypes]", tools) if isinstance(tools, list) else None
approval_tools = [tool_name for tool_name, tool in tool_map.items() if tool.approval_mode == "always_require"]
approval_tools = {tool_name for tool_name, tool in tool_map.items() if tool.approval_mode == "always_require"}
logger.debug(
"_try_execute_function_calls: tool_map keys=%s, approval_tools=%s",
list(tool_map.keys()),
approval_tools,
)
declaration_only = [tool_name for tool_name, tool in tool_map.items() if tool.declaration_only]
declaration_only = {tool_name for tool_name, tool in tool_map.items() if tool.declaration_only}
configured_additional_tools = config.get("additional_tools") or []
additional_tool_names = [tool.name for tool in configured_additional_tools]
additional_tool_names = {tool.name for tool in configured_additional_tools}
# check if any are calling functions that need approval
# if so, we return approval request for all
approval_needed = False
@@ -1719,15 +1722,39 @@ async def _try_execute_function_calls(
raise KeyError(f'Error: Requested function "{fcc.name}" not found.') # type: ignore[attr-defined]
if approval_needed:
# approval can only be needed for Function Call Content, not Approval Responses.
logger.debug("Returning function_approval_request contents")
return (
[
Content.from_function_approval_request(id=fcc.call_id, function_call=fcc) # type: ignore[attr-defined, arg-type]
for fcc in function_calls
if fcc.type == "function_call"
],
False,
logger.debug("Returning visible function_approval_request contents and storing already-approved requests")
visible_requests: list[Content] = []
already_approved_requests: list[Content] = []
for fcc in function_calls:
if fcc.type != "function_call":
continue
approval_request = Content.from_function_approval_request(
id=fcc.call_id, # type: ignore[arg-type]
function_call=fcc,
)
tool_name = fcc.name # type: ignore[attr-defined]
if tool_name is None:
visible_requests.append(approval_request)
continue
tool = tool_map.get(tool_name)
if (
tool_name in approval_tools
or tool is None
or tool_name in declaration_only
or tool_name in additional_tool_names
):
visible_requests.append(approval_request)
continue
if invocation_session is None:
visible_requests.append(approval_request)
continue
already_approved_requests.append(approval_request)
_store_already_approved_approval_requests(
invocation_session,
visible_requests,
already_approved_requests,
)
return (visible_requests, False)
if declaration_only_flag:
# return the declaration only tools to the user, since we cannot execute them.
# Mark as user_input_request so AgentExecutor emits request_info events and pauses the workflow.
@@ -1912,6 +1939,108 @@ def _is_hosted_tool_approval(content: Any) -> bool:
return bool(ap and ap.get("server_label"))
def _get_tool_approval_state(invocation_session: AgentSession | None) -> dict[str, Any] | None:
"""Return the shared tool-approval state bag for the invocation session."""
if invocation_session is None:
return None
raw_state = invocation_session.state.get(_TOOL_APPROVAL_STATE_KEY)
if isinstance(raw_state, dict):
return cast(dict[str, Any], raw_state)
from ._harness._tool_approval import ToolApprovalState
if isinstance(raw_state, ToolApprovalState):
serialized_state = raw_state.to_dict(exclude={"type"})
invocation_session.state[_TOOL_APPROVAL_STATE_KEY] = serialized_state
return serialized_state
if raw_state is not None:
raise TypeError(
f"Session state for {_TOOL_APPROVAL_STATE_KEY!r} must be a dict or ToolApprovalState, "
f"got {type(raw_state).__name__}."
)
new_state: dict[str, Any] = {}
invocation_session.state[_TOOL_APPROVAL_STATE_KEY] = new_state
return new_state
def _content_from_state(value: Any) -> Content | None:
"""Restore a Content item stored in session state."""
from ._types import Content
if isinstance(value, Content):
return value
if isinstance(value, Mapping):
return Content.from_dict(cast(Mapping[str, Any], value))
return None
def _store_already_approved_approval_requests(
invocation_session: AgentSession | None,
visible_approval_requests: Sequence[Content],
already_approved_requests: Sequence[Content],
) -> None:
"""Store hidden already-approved requests keyed by the visible approvals that resume the batch."""
if not already_approved_requests:
return
state = _get_tool_approval_state(invocation_session)
if state is None:
return
visible_ids = [request.id for request in visible_approval_requests if request.id]
if not visible_ids:
return
existing_groups = state.get(_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY)
pending_groups: list[Any] = (
list(cast(Iterable[Any], existing_groups)) if isinstance(existing_groups, list) else []
)
pending_groups.append({
"approval_request_ids": visible_ids,
"approval_requests": [request.to_dict() for request in already_approved_requests],
})
state[_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY] = pending_groups
def _pop_already_approved_approval_responses(
invocation_session: AgentSession | None,
approval_response_ids: set[str],
) -> list[Content]:
"""Pop already-approved requests for the visible approval ids being answered."""
if not approval_response_ids:
return []
state = _get_tool_approval_state(invocation_session)
if state is None:
return []
raw_groups = state.get(_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY, [])
if not isinstance(raw_groups, list):
return []
responses: list[Content] = []
remaining_groups: list[Any] = []
raw_group_items = list(cast(Iterable[Any], raw_groups))
for raw_group in raw_group_items:
if not isinstance(raw_group, Mapping):
continue
group = cast(Mapping[str, Any], raw_group)
raw_ids = group.get("approval_request_ids")
raw_group_ids: Iterable[Any] = cast(Iterable[Any], raw_ids) if isinstance(raw_ids, list) else ()
group_ids = {str(item) for item in raw_group_ids}
if group_ids.isdisjoint(approval_response_ids):
remaining_groups.append(raw_group)
continue
raw_requests = group.get("approval_requests")
if not isinstance(raw_requests, list):
continue
for raw_request in list(cast(Iterable[Any], raw_requests)):
request = _content_from_state(raw_request)
if request is None or request.type != "function_approval_request":
continue
responses.append(request.to_function_approval_response(approved=True))
if remaining_groups:
state[_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY] = remaining_groups
else:
state.pop(_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY, None)
return responses
def _collect_approval_responses(
messages: list[Message],
) -> dict[str, Content]:
@@ -2157,8 +2286,24 @@ async def _process_function_requests(
errors_in_a_row: int,
max_errors: int,
execute_function_calls: Callable[..., Awaitable[tuple[list[Content], bool, bool]]],
invocation_session: AgentSession | None = None,
) -> FunctionRequestResult:
from ._types import Message
if prepped_messages is not None:
explicit_approval_response_ids = {
content.id
for message in prepped_messages
if isinstance(message, Message)
for content in message.contents
if content.type == "function_approval_response" and content.id
}
already_approved_responses = _pop_already_approved_approval_responses(
invocation_session,
explicit_approval_response_ids,
)
if already_approved_responses:
prepped_messages.append(Message(role="user", contents=already_approved_responses))
fcc_todo = _collect_approval_responses(prepped_messages)
if not fcc_todo:
fcc_todo = {}
@@ -2362,6 +2507,10 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
function_middleware_pipeline = self._get_function_middleware_pipeline(runtime_middleware["function"])
if runtime_middleware["chat"]:
effective_client_kwargs["middleware"] = runtime_middleware["chat"]
raw_budget_state = effective_client_kwargs.pop(_FUNCTION_INVOCATION_BUDGET_STATE_KEY, None)
budget_state: dict[str, Any] = (
cast(dict[str, Any], raw_budget_state) if isinstance(raw_budget_state, dict) else {}
)
max_errors = self.function_invocation_configuration.get(
"max_consecutive_errors_per_request", DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST
)
@@ -2411,7 +2560,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
nonlocal mutable_options
nonlocal filtered_kwargs
errors_in_a_row: int = 0
total_function_calls: int = 0
total_function_calls = int(budget_state.get("total_function_calls", 0) or 0)
max_function_calls: int | None = self.function_invocation_configuration.get("max_function_calls")
prepped_messages = list(messages)
fcc_messages: list[Message] = []
@@ -2420,7 +2569,9 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
loop_enabled = self.function_invocation_configuration.get("enabled", True)
max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS)
for attempt_idx in range(max_iterations if loop_enabled else 0):
attempt_start = int(budget_state.get("attempt_count", 0) or 0)
for attempt_idx in range(attempt_start, max_iterations if loop_enabled else 0):
budget_state["attempt_count"] = attempt_idx + 1
approval_result = await _process_function_requests(
response=None,
prepped_messages=prepped_messages,
@@ -2430,12 +2581,21 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
errors_in_a_row=errors_in_a_row,
max_errors=max_errors,
execute_function_calls=execute_function_calls,
invocation_session=invocation_session,
)
if approval_result.get("action") == "stop":
response = ChatResponse(messages=prepped_messages)
break
errors_in_a_row = approval_result.get("errors_in_a_row", errors_in_a_row)
total_function_calls += approval_result.get("function_call_count", 0)
budget_state["total_function_calls"] = total_function_calls
if max_function_calls is not None and total_function_calls >= max_function_calls:
logger.info(
"Maximum function calls reached (%d/%d). Stopping further function calls for this request.",
total_function_calls,
max_function_calls,
)
mutable_options["tool_choice"] = "none"
response = cast(
ChatResponse[Any],
@@ -2468,11 +2628,13 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
errors_in_a_row=errors_in_a_row,
max_errors=max_errors,
execute_function_calls=execute_function_calls,
invocation_session=invocation_session,
)
if result.get("action") == "return":
response.usage_details = aggregated_usage
return _clear_internal_conversation_id(response)
total_function_calls += result.get("function_call_count", 0)
budget_state["total_function_calls"] = total_function_calls
if result.get("action") == "stop":
# Error threshold reached: force a final non-tool turn so
# function_call_output items are submitted before exit.
@@ -2549,7 +2711,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
nonlocal mutable_options
nonlocal stream_result_hooks
errors_in_a_row: int = 0
total_function_calls: int = 0
total_function_calls = int(budget_state.get("total_function_calls", 0) or 0)
max_function_calls: int | None = self.function_invocation_configuration.get("max_function_calls")
prepped_messages = list(messages)
fcc_messages: list[Message] = []
@@ -2557,7 +2719,9 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
loop_enabled = self.function_invocation_configuration.get("enabled", True)
max_iterations = self.function_invocation_configuration.get("max_iterations", DEFAULT_MAX_ITERATIONS)
for attempt_idx in range(max_iterations if loop_enabled else 0):
attempt_start = int(budget_state.get("attempt_count", 0) or 0)
for attempt_idx in range(attempt_start, max_iterations if loop_enabled else 0):
budget_state["attempt_count"] = attempt_idx + 1
approval_result = await _process_function_requests(
response=None,
prepped_messages=prepped_messages,
@@ -2567,9 +2731,18 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
errors_in_a_row=errors_in_a_row,
max_errors=max_errors,
execute_function_calls=execute_function_calls,
invocation_session=invocation_session,
)
errors_in_a_row = approval_result.get("errors_in_a_row", errors_in_a_row)
total_function_calls += approval_result.get("function_call_count", 0)
budget_state["total_function_calls"] = total_function_calls
if max_function_calls is not None and total_function_calls >= max_function_calls:
logger.info(
"Maximum function calls reached (%d/%d). Stopping further function calls for this request.",
total_function_calls,
max_function_calls,
)
mutable_options["tool_choice"] = "none"
if approval_result.get("action") == "stop":
mutable_options["tool_choice"] = "none"
return
@@ -2622,9 +2795,11 @@ class FunctionInvocationLayer(Generic[OptionsCoT]):
errors_in_a_row=errors_in_a_row,
max_errors=max_errors,
execute_function_calls=execute_function_calls,
invocation_session=invocation_session,
)
errors_in_a_row = result.get("errors_in_a_row", errors_in_a_row)
total_function_calls += result.get("function_call_count", 0)
budget_state["total_function_calls"] = total_function_calls
if role := result.get("update_role"):
yield ChatResponseUpdate(
contents=result.get("function_call_results") or [],
python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py
+29
View File
@@ -13,6 +13,35 @@ during deserialization. The default built-in safe set covers common Python
value types (primitives, datetime, uuid, ...), all ``agent_framework`` internal
types, and all ``openai.types`` types. Callers can extend the set by passing
additional ``"module:qualname"`` strings.
Security Model
--------------
Checkpoint storage is treated as a **trusted data source**. The serialization
format uses Python's ``pickle`` module which can execute arbitrary code during
deserialization. The ``RestrictedUnpickler`` provides a defense-in-depth
allowlist that limits instantiable classes, but it is **not** a security
boundary — certain allowlisted builtins (e.g. ``getattr``) are required for
legitimate object reconstruction (enums, named tuples) and cannot be removed
without breaking compatibility.
Developers **must** ensure that:
1. The checkpoint storage backend (file system, Cosmos DB, Azure Blob, Durable
Functions storage) is access-controlled and not writable by untrusted
parties.
2. Data flowing into ``decode_checkpoint_value`` originates exclusively from
the application's own checkpoint storage — never from user-supplied HTTP
requests, message payloads, or other untrusted sources.
3. The ``allowed_types`` parameter is specified whenever possible to restrict
the set of reconstructible types to the minimum required by the application.
4. Never pass untrusted external input to ``decode_checkpoint_value``. If you
must accept external JSON that might contain checkpoint markers, sanitize it
first (for example, :func:`agent_framework_azurefunctions._serialization.strip_pickle_markers`).
The allowlist is a mitigation that reduces attack surface but does not
eliminate the inherent risks of deserializing untrusted pickle data. Treat
your checkpoint storage with the same access controls you would apply to
application secrets or database credentials.
"""
from __future__ import annotations
python/packages/core/pyproject.toml
+1 -1
View File
@@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.8.0"
version = "1.8.1"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
python/packages/core/tests/core/test_harness_agent.py
+190 -8
View File
@@ -194,6 +194,63 @@ def test_create_harness_agent_returns_full_agent() -> None:
assert isinstance(agent, FullAgent)
def test_create_harness_agent_no_token_params_disables_compaction() -> None:
"""When token params are omitted, compaction is automatically disabled."""
agent = create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
)
provider_types = [type(p) for p in agent.context_providers]
assert CompactionProvider not in provider_types
def test_create_harness_agent_no_token_params_skips_max_tokens_option() -> None:
"""When max_output_tokens is omitted, max_tokens should not be set in default options."""
agent = create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
)
assert agent.default_options.get("max_tokens") is None
def test_create_harness_agent_custom_before_strategy_enables_compaction_without_tokens() -> None:
"""A custom before_compaction_strategy enables compaction even when token params are omitted."""
from agent_framework import ToolResultCompactionStrategy
agent = create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
before_compaction_strategy=ToolResultCompactionStrategy(),
)
provider_types = [type(p) for p in agent.context_providers]
assert CompactionProvider in provider_types
def test_create_harness_agent_disable_compaction_overrides_custom_before_strategy() -> None:
"""disable_compaction=True wins even when a custom before strategy is provided."""
from agent_framework import ToolResultCompactionStrategy
agent = create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
before_compaction_strategy=ToolResultCompactionStrategy(),
disable_compaction=True,
)
provider_types = [type(p) for p in agent.context_providers]
assert CompactionProvider not in provider_types
def test_create_harness_agent_custom_after_strategy_enables_compaction_without_tokens() -> None:
"""A custom after_compaction_strategy enables compaction even when token params are omitted."""
from agent_framework import ToolResultCompactionStrategy
agent = create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
after_compaction_strategy=ToolResultCompactionStrategy(),
)
compaction_providers = [p for p in agent.context_providers if isinstance(p, CompactionProvider)]
assert len(compaction_providers) == 1
# Before phase is skipped (no token budget, no custom before strategy), after phase is set.
assert compaction_providers[0].before_strategy is None
assert compaction_providers[0].after_strategy is not None
# --- Validation Tests ---
@@ -207,14 +264,15 @@ def test_create_harness_agent_rejects_invalid_context_tokens() -> None:
)
def test_create_harness_agent_rejects_negative_output_tokens() -> None:
"""max_output_tokens must be non-negative."""
with pytest.raises(ValueError, match="max_output_tokens must be non-negative"):
create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
max_context_window_tokens=1000,
max_output_tokens=-1,
)
def test_create_harness_agent_rejects_non_positive_output_tokens() -> None:
"""max_output_tokens must be positive when provided."""
for invalid_value in (0, -1):
with pytest.raises(ValueError, match="max_output_tokens must be positive"):
create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
max_context_window_tokens=1000,
max_output_tokens=invalid_value,
)
def test_create_harness_agent_rejects_output_gte_context() -> None:
@@ -485,3 +543,127 @@ def test_create_harness_agent_empty_background_agents_list() -> None:
)
providers = agent.context_providers or []
assert not any(isinstance(p, BackgroundAgentsProvider) for p in providers)
# --- Shell Tool Tests ---
class _FakeShellTool:
"""Fake shell executor/tool exposing as_function()."""
def as_function(self) -> str:
return "shell_fn"
class _FakeShellClient(_FakeChatClient):
"""Fake client that supports the shell tool."""
def __init__(self) -> None:
self.shell_func: Any = None
def get_shell_tool(self, *, func: Any = None, **kwargs: Any) -> str:
self.shell_func = func
return "shell_tool_instance"
def test_create_harness_agent_adds_shell_tool_and_provider() -> None:
"""Shell tool and ShellEnvironmentProvider should be added when a shell executor is supplied."""
from agent_framework_tools.shell import ShellEnvironmentProvider
client = _FakeShellClient()
agent = create_harness_agent(
client=client, # type: ignore[arg-type]
max_context_window_tokens=128_000,
max_output_tokens=16_384,
disable_web_search=True,
shell_executor=_FakeShellTool(),
)
tools = agent.default_options.get("tools", [])
assert "shell_tool_instance" in tools
assert client.shell_func == "shell_fn"
providers = agent.context_providers or []
assert any(isinstance(p, ShellEnvironmentProvider) for p in providers)
def test_create_harness_agent_shell_passes_custom_options() -> None:
"""Custom ShellEnvironmentProviderOptions should be forwarded to the provider."""
from agent_framework_tools.shell import ShellEnvironmentProvider, ShellEnvironmentProviderOptions
options = ShellEnvironmentProviderOptions(probe_tools=("git",))
agent = create_harness_agent(
client=_FakeShellClient(), # type: ignore[arg-type]
max_context_window_tokens=128_000,
max_output_tokens=16_384,
disable_web_search=True,
shell_executor=_FakeShellTool(),
shell_environment_provider_options=options,
)
providers = agent.context_providers or []
provider = next(p for p in providers if isinstance(p, ShellEnvironmentProvider))
assert provider._options is options
def test_create_harness_agent_shell_skipped_when_unsupported(caplog: pytest.LogCaptureFixture) -> None:
"""When the client lacks get_shell_tool, both the tool and provider are skipped with a warning."""
import logging
from agent_framework_tools.shell import ShellEnvironmentProvider
with caplog.at_level(logging.WARNING, logger="agent_framework._harness._agent"):
agent = create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
max_context_window_tokens=128_000,
max_output_tokens=16_384,
disable_web_search=True,
shell_executor=_FakeShellTool(),
)
assert any("SupportsShellTool" in msg for msg in caplog.messages)
providers = agent.context_providers or []
assert not any(isinstance(p, ShellEnvironmentProvider) for p in providers)
assert "tools" not in agent.default_options or not agent.default_options.get("tools")
def test_create_harness_agent_no_shell_by_default() -> None:
"""No shell tool or provider should be added when shell_executor is not provided."""
from agent_framework_tools.shell import ShellEnvironmentProvider
agent = create_harness_agent(
client=_FakeShellClient(), # type: ignore[arg-type]
max_context_window_tokens=128_000,
max_output_tokens=16_384,
disable_web_search=True,
)
providers = agent.context_providers or []
assert not any(isinstance(p, ShellEnvironmentProvider) for p in providers)
def test_create_harness_agent_shell_executor_without_as_function_raises() -> None:
"""A shell_executor lacking a callable as_function() should raise a clear TypeError."""
class _BadExecutor:
pass
with pytest.raises(TypeError, match="as_function"):
create_harness_agent(
client=_FakeShellClient(), # type: ignore[arg-type]
max_context_window_tokens=128_000,
max_output_tokens=16_384,
disable_web_search=True,
shell_executor=_BadExecutor(),
)
def test_create_harness_agent_shell_executor_validated_before_client_check() -> None:
"""The as_function() contract is validated upfront, even when the client lacks shell support."""
class _BadExecutor:
pass
with pytest.raises(TypeError, match="as_function"):
create_harness_agent(
client=_FakeChatClient(), # type: ignore[arg-type]
max_context_window_tokens=128_000,
max_output_tokens=16_384,
disable_web_search=True,
shell_executor=_BadExecutor(),
)
python/packages/core/tests/core/test_harness_tool_approval.py
+817
View File
@@ -0,0 +1,817 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from agent_framework import (
DEFAULT_TOOL_APPROVAL_SOURCE_ID,
Agent,
AgentSession,
ChatResponse,
ChatResponseUpdate,
Content,
Message,
SupportsChatGetResponse,
ToolApprovalMiddleware,
ToolApprovalState,
create_always_approve_tool_response,
create_always_approve_tool_with_arguments_response,
tool,
)
def _approval_requests(messages: list[Message]) -> list[Content]:
return [
content for message in messages for content in message.contents if content.type == "function_approval_request"
]
async def test_mixed_batch_hides_already_approved_request_until_approval_replay(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Mixed batches should only show real approval requests when a session can store hidden requests."""
no_approval_calls = 0
approval_calls = 0
@tool(name="lookup_work_items", approval_mode="never_require")
def lookup_work_items(query: str) -> str:
nonlocal no_approval_calls
no_approval_calls += 1
return f"found {query}"
@tool(name="add_comment", approval_mode="always_require")
def add_comment(comment: str) -> str:
nonlocal approval_calls
approval_calls += 1
return f"added {comment}"
agent = Agent(client=chat_client_base, tools=[lookup_work_items, add_comment])
session = AgentSession(session_id="approval-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_lookup",
name="lookup_work_items",
arguments='{"query": "mine"}',
),
Content.from_function_call(
call_id="call_comment",
name="add_comment",
arguments='{"comment": "done"}',
),
],
)
)
]
first_response = await agent.run("update work item", session=session)
requests = _approval_requests(first_response.messages)
assert [request.function_call.name for request in requests] == ["add_comment"]
assert no_approval_calls == 0
assert approval_calls == 0
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["complete"]))]
second_response = await agent.run(requests[0].to_function_approval_response(approved=True), session=session)
assert second_response.text == "complete"
assert no_approval_calls == 1
assert approval_calls == 1
async def test_mixed_batch_accepts_restored_tool_approval_state(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Mixed-batch bypass should work when session state contains ToolApprovalState."""
safe_calls = 0
risky_calls = 0
@tool(name="safe_read", approval_mode="never_require")
def safe_read() -> str:
nonlocal safe_calls
safe_calls += 1
return "safe"
@tool(name="risky_write", approval_mode="always_require")
def risky_write() -> str:
nonlocal risky_calls
risky_calls += 1
return "risky"
agent = Agent(client=chat_client_base, tools=[safe_read, risky_write])
session = AgentSession(session_id="restored-state-session")
session.state[DEFAULT_TOOL_APPROVAL_SOURCE_ID] = ToolApprovalState()
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_safe", name="safe_read", arguments="{}"),
Content.from_function_call(call_id="call_risky", name="risky_write", arguments="{}"),
],
)
)
]
first_response = await agent.run("read and write", session=session)
requests = _approval_requests(first_response.messages)
assert [request.function_call.name for request in requests] == ["risky_write"]
assert safe_calls == 0
assert risky_calls == 0
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
final_response = await agent.run(requests[0].to_function_approval_response(approved=True), session=session)
assert final_response.text == "done"
assert safe_calls == 1
assert risky_calls == 1
async def test_hidden_mixed_batch_requests_do_not_replay_on_unrelated_turn(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Stored hidden approvals should only replay when an approval response resumes the flow."""
safe_calls = 0
risky_calls = 0
@tool(name="safe_lookup", approval_mode="never_require")
def safe_lookup() -> str:
nonlocal safe_calls
safe_calls += 1
return "safe"
@tool(name="risky_update", approval_mode="always_require")
def risky_update() -> str:
nonlocal risky_calls
risky_calls += 1
return "risky"
agent = Agent(client=chat_client_base, tools=[safe_lookup, risky_update])
session = AgentSession(session_id="stale-hidden-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_safe", name="safe_lookup", arguments="{}"),
Content.from_function_call(call_id="call_risky", name="risky_update", arguments="{}"),
],
)
)
]
first_response = await agent.run("lookup and update", session=session)
request = _approval_requests(first_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["unrelated"]))]
unrelated_response = await agent.run("never mind, answer something else", session=session)
assert unrelated_response.text == "unrelated"
assert safe_calls == 0
assert risky_calls == 0
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
final_response = await agent.run(request.to_function_approval_response(approved=True), session=session)
assert final_response.text == "done"
assert safe_calls == 1
assert risky_calls == 1
async def test_hidden_mixed_batch_requests_replay_only_for_matching_visible_approval(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Approving one mixed batch must not replay hidden calls from another abandoned batch."""
safe_a_calls = 0
safe_b_calls = 0
risky_a_calls = 0
risky_b_calls = 0
@tool(name="safe_a", approval_mode="never_require")
def safe_a() -> str:
nonlocal safe_a_calls
safe_a_calls += 1
return "safe-a"
@tool(name="safe_b", approval_mode="never_require")
def safe_b() -> str:
nonlocal safe_b_calls
safe_b_calls += 1
return "safe-b"
@tool(name="risky_a", approval_mode="always_require")
def risky_a() -> str:
nonlocal risky_a_calls
risky_a_calls += 1
return "risky-a"
@tool(name="risky_b", approval_mode="always_require")
def risky_b() -> str:
nonlocal risky_b_calls
risky_b_calls += 1
return "risky-b"
agent = Agent(client=chat_client_base, tools=[safe_a, safe_b, risky_a, risky_b])
session = AgentSession(session_id="grouped-hidden-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_safe_a", name="safe_a", arguments="{}"),
Content.from_function_call(call_id="call_risky_a", name="risky_a", arguments="{}"),
],
)
)
]
first_response = await agent.run("batch a", session=session)
assert [request.function_call.name for request in _approval_requests(first_response.messages)] == ["risky_a"]
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_safe_b", name="safe_b", arguments="{}"),
Content.from_function_call(call_id="call_risky_b", name="risky_b", arguments="{}"),
],
)
)
]
second_response = await agent.run("batch b", session=session)
second_request = _approval_requests(second_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
final_response = await agent.run(second_request.to_function_approval_response(approved=True), session=session)
assert final_response.text == "done"
assert safe_a_calls == 0
assert risky_a_calls == 0
assert safe_b_calls == 1
assert risky_b_calls == 1
async def test_tool_approval_middleware_queues_multiple_approval_requests(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""The opt-in middleware should present multiple unresolved approvals one at a time."""
first_calls = 0
second_calls = 0
@tool(name="first_tool", approval_mode="always_require")
def first_tool() -> str:
nonlocal first_calls
first_calls += 1
return "first"
@tool(name="second_tool", approval_mode="always_require")
def second_tool() -> str:
nonlocal second_calls
second_calls += 1
return "second"
agent = Agent(
client=chat_client_base,
tools=[first_tool, second_tool],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="queue-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_first", name="first_tool", arguments="{}"),
Content.from_function_call(call_id="call_second", name="second_tool", arguments="{}"),
],
)
)
]
first_response = await agent.run("call both", session=session)
first_requests = _approval_requests(first_response.messages)
assert [request.function_call.name for request in first_requests] == ["first_tool"]
assert first_calls == 0
assert second_calls == 0
second_response = await agent.run(first_requests[0].to_function_approval_response(approved=True), session=session)
second_requests = _approval_requests(second_response.messages)
assert [request.function_call.name for request in second_requests] == ["second_tool"]
assert first_calls == 0
assert second_calls == 0
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
final_response = await agent.run(second_requests[0].to_function_approval_response(approved=True), session=session)
assert final_response.text == "done"
assert first_calls == 1
assert second_calls == 1
async def test_tool_approval_middleware_preserves_hidden_mixed_batch_requests(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Middleware state saves should not discard core hidden already-approved requests."""
lookup_calls = 0
write_calls = 0
@tool(name="lookup_records", approval_mode="never_require")
def lookup_records() -> str:
nonlocal lookup_calls
lookup_calls += 1
return "records"
@tool(name="write_record", approval_mode="always_require")
def write_record() -> str:
nonlocal write_calls
write_calls += 1
return "written"
agent = Agent(
client=chat_client_base,
tools=[lookup_records, write_record],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="mixed-middleware-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_lookup", name="lookup_records", arguments="{}"),
Content.from_function_call(call_id="call_write", name="write_record", arguments="{}"),
],
)
)
]
first_response = await agent.run("lookup and write", session=session)
request = _approval_requests(first_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
second_response = await agent.run(request.to_function_approval_response(approved=True), session=session)
assert second_response.text == "done"
assert lookup_calls == 1
assert write_calls == 1
async def test_tool_approval_middleware_auto_approval_rule_receives_function_call(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Heuristic auto-approval callbacks should receive function-call content and approve matching calls."""
auto_calls = 0
manual_calls = 0
seen_calls: list[tuple[str, str | None]] = []
@tool(name="auto_write", approval_mode="always_require")
def auto_write() -> str:
nonlocal auto_calls
auto_calls += 1
return "auto"
@tool(name="manual_write", approval_mode="always_require")
def manual_write() -> str:
nonlocal manual_calls
manual_calls += 1
return "manual"
async def auto_approve_auto_write(function_call: Content) -> bool:
seen_calls.append((function_call.type, function_call.name))
return function_call.name == "auto_write"
agent = Agent(
client=chat_client_base,
tools=[auto_write, manual_write],
middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_auto_write])],
)
session = AgentSession(session_id="heuristic-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(call_id="call_auto", name="auto_write", arguments="{}"),
Content.from_function_call(call_id="call_manual", name="manual_write", arguments="{}"),
],
)
)
]
first_response = await agent.run("write both", session=session)
requests = _approval_requests(first_response.messages)
assert [request.function_call.name for request in requests] == ["manual_write"]
assert seen_calls == [("function_call", "auto_write"), ("function_call", "manual_write")]
assert auto_calls == 0
assert manual_calls == 0
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["done"]))]
final_response = await agent.run(requests[0].to_function_approval_response(approved=True), session=session)
assert final_response.text == "done"
assert auto_calls == 1
assert manual_calls == 1
async def test_tool_approval_middleware_auto_approved_loops_share_function_call_budget(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Auto-approved re-entry should not reset max_function_calls."""
calls = 0
@tool(name="budgeted_tool", approval_mode="always_require")
def budgeted_tool(value: str) -> str:
nonlocal calls
calls += 1
return value
def auto_approve_budgeted_tool(function_call: Content) -> bool:
return function_call.name == "budgeted_tool"
chat_client_base.function_invocation_configuration["max_function_calls"] = 1 # type: ignore[attr-defined]
agent = Agent(
client=chat_client_base,
tools=[budgeted_tool],
middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_budgeted_tool])],
)
session = AgentSession(session_id="shared-budget-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_first",
name="budgeted_tool",
arguments='{"value": "first"}',
)
],
)
),
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_second",
name="budgeted_tool",
arguments='{"value": "second"}',
)
],
)
),
]
response = await agent.run("call repeatedly", session=session)
assert response.text == "I broke out of the function invocation loop..."
assert calls == 1
async def test_tool_approval_middleware_queues_streamed_approval_requests(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Streaming approval requests should also be queued one at a time."""
calls = 0
@tool(name="first_streamed_tool", approval_mode="always_require")
def first_streamed_tool() -> str:
nonlocal calls
calls += 1
return "first"
@tool(name="second_streamed_tool", approval_mode="always_require")
def second_streamed_tool() -> str:
nonlocal calls
calls += 1
return "second"
agent = Agent(
client=chat_client_base,
tools=[first_streamed_tool, second_streamed_tool],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="stream-queue-session")
chat_client_base.streaming_responses = [
[
ChatResponseUpdate(
contents=[Content.from_function_call(call_id="call_first", name="first_streamed_tool", arguments="{}")],
role="assistant",
),
ChatResponseUpdate(
contents=[
Content.from_function_call(call_id="call_second", name="second_streamed_tool", arguments="{}")
],
role="assistant",
),
]
]
first_stream = agent.run("call both", stream=True, session=session)
first_updates = [update async for update in first_stream]
first_requests = [content for update in first_updates for content in update.user_input_requests]
assert [request.function_call.name for request in first_requests] == ["first_streamed_tool"]
assert calls == 0
second_stream = agent.run(
first_requests[0].to_function_approval_response(approved=True),
stream=True,
session=session,
)
second_updates = [update async for update in second_stream]
second_requests = [content for update in second_updates for content in update.user_input_requests]
assert [request.function_call.name for request in second_requests] == ["second_streamed_tool"]
assert calls == 0
chat_client_base.streaming_responses = [
[ChatResponseUpdate(contents=[Content.from_text("done")], role="assistant")]
]
final_stream = agent.run(
second_requests[0].to_function_approval_response(approved=True),
stream=True,
session=session,
)
final_updates = [update async for update in final_stream]
final_response = await final_stream.get_final_response()
assert final_updates[-1].text == "done"
assert final_response.text == "done"
assert calls == 2
async def test_tool_approval_middleware_always_approve_tool_rule(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""An always-approve response should add a standing tool-level approval rule."""
calls = 0
@tool(name="dangerous_tool", approval_mode="always_require")
def dangerous_tool(value: str) -> str:
nonlocal calls
calls += 1
return value
agent = Agent(
client=chat_client_base,
tools=[dangerous_tool],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="standing-rule-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_initial",
name="dangerous_tool",
arguments='{"value": "one"}',
)
],
)
)
]
first_response = await agent.run("call once", session=session)
first_request = _approval_requests(first_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["first done"]))]
await agent.run(create_always_approve_tool_response(first_request), session=session)
assert calls == 1
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_auto",
name="dangerous_tool",
arguments='{"value": "two"}',
)
],
)
),
ChatResponse(messages=Message(role="assistant", contents=["second done"])),
]
second_response = await agent.run("call again", session=session)
assert second_response.text == "second done"
assert calls == 2
async def test_tool_approval_middleware_standing_rules_include_hosted_server_boundary(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""A standing hosted-tool rule should only match the same server_label."""
calls = 0
@tool(name="hosted_tool", approval_mode="always_require")
def hosted_tool() -> str:
nonlocal calls
calls += 1
return "hosted"
def hosted_call(call_id: str, server_label: str) -> Content:
return Content.from_function_call(
call_id=call_id,
name="hosted_tool",
arguments="{}",
additional_properties={"server_label": server_label},
)
agent = Agent(
client=chat_client_base,
tools=[hosted_tool],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="hosted-boundary-session")
chat_client_base.run_responses = [
ChatResponse(messages=Message(role="assistant", contents=[hosted_call("call_initial", "server-a")]))
]
first_response = await agent.run("call hosted a", session=session)
first_request = _approval_requests(first_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["server a done"]))]
await agent.run(create_always_approve_tool_response(first_request), session=session)
assert calls == 0
chat_client_base.run_responses = [
ChatResponse(messages=Message(role="assistant", contents=[hosted_call("call_same_server", "server-a")])),
ChatResponse(messages=Message(role="assistant", contents=["same server done"])),
]
same_server_response = await agent.run("call hosted a again", session=session)
assert same_server_response.text == "same server done"
assert _approval_requests(same_server_response.messages) == []
assert calls == 0
chat_client_base.run_responses = [
ChatResponse(messages=Message(role="assistant", contents=[hosted_call("call_other_server", "server-b")]))
]
other_server_response = await agent.run("call hosted b", session=session)
requests = _approval_requests(other_server_response.messages)
assert [request.function_call.additional_properties["server_label"] for request in requests] == ["server-b"]
assert calls == 0
async def test_tool_approval_middleware_always_approve_tool_with_arguments_rule(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""Argument-scoped always-approve rules should require exact argument matches."""
calls = 0
@tool(name="argument_scoped_tool", approval_mode="always_require")
def argument_scoped_tool(value: str) -> str:
nonlocal calls
calls += 1
return value
agent = Agent(
client=chat_client_base,
tools=[argument_scoped_tool],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="argument-rule-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_initial",
name="argument_scoped_tool",
arguments='{"value": "same"}',
)
],
)
)
]
first_response = await agent.run("call with same", session=session)
first_request = _approval_requests(first_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["first done"]))]
await agent.run(create_always_approve_tool_with_arguments_response(first_request), session=session)
assert calls == 1
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_same",
name="argument_scoped_tool",
arguments='{"value": "same"}',
)
],
)
),
ChatResponse(messages=Message(role="assistant", contents=["same done"])),
]
second_response = await agent.run("call with same again", session=session)
assert second_response.text == "same done"
assert calls == 2
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_different",
name="argument_scoped_tool",
arguments='{"value": "different"}',
)
],
)
)
]
third_response = await agent.run("call with different args", session=session)
requests = _approval_requests(third_response.messages)
assert [request.function_call.arguments for request in requests] == ['{"value": "different"}']
assert calls == 2
async def test_tool_approval_middleware_empty_arguments_rule_is_not_tool_wide(
chat_client_base: SupportsChatGetResponse,
) -> None:
"""An argument-scoped no-argument approval should not become a wildcard."""
calls = 0
@tool(name="optional_args_tool", approval_mode="always_require")
def optional_args_tool(value: str = "default") -> str:
nonlocal calls
calls += 1
return value
agent = Agent(
client=chat_client_base,
tools=[optional_args_tool],
middleware=[ToolApprovalMiddleware()],
)
session = AgentSession(session_id="empty-arguments-rule-session")
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_empty",
name="optional_args_tool",
arguments="{}",
)
],
)
)
]
first_response = await agent.run("call without args", session=session)
first_request = _approval_requests(first_response.messages)[0]
chat_client_base.run_responses = [ChatResponse(messages=Message(role="assistant", contents=["empty done"]))]
await agent.run(create_always_approve_tool_with_arguments_response(first_request), session=session)
assert calls == 1
chat_client_base.run_responses = [
ChatResponse(
messages=Message(
role="assistant",
contents=[
Content.from_function_call(
call_id="call_non_empty",
name="optional_args_tool",
arguments='{"value": "custom"}',
)
],
)
)
]
second_response = await agent.run("call with args", session=session)
requests = _approval_requests(second_response.messages)
assert [request.function_call.arguments for request in requests] == ['{"value": "custom"}']
assert calls == 1
python/packages/core/tests/core/test_mcp.py
+277 -20
View File
@@ -342,6 +342,69 @@ def test_parse_tool_result_from_mcp_resource_link_text_resource_and_unknown():
assert result[1].text == "Embedded result"
def test_parse_tool_result_from_mcp_structured_content_only():
"""Test that structuredContent is parsed when content list is empty."""
mcp_result = types.CallToolResult(
content=[],
structuredContent={"Tables": [{"Name": "Sales", "Columns": ["Amount", "Date"]}]},
)
result = _HELPER_MCP_TOOL._parse_tool_result_from_mcp(mcp_result)
assert isinstance(result, list)
assert len(result) == 1
assert result[0].type == "text"
parsed = json.loads(result[0].text)
assert parsed == {"Tables": [{"Name": "Sales", "Columns": ["Amount", "Date"]}]}
def test_parse_tool_result_from_mcp_structured_content_with_text():
"""Test that structuredContent is appended alongside regular content items."""
mcp_result = types.CallToolResult(
content=[types.TextContent(type="text", text="Summary")],
structuredContent={"data": [1, 2, 3]},
)
result = _HELPER_MCP_TOOL._parse_tool_result_from_mcp(mcp_result)
assert isinstance(result, list)
assert len(result) == 2
assert result[0].type == "text"
assert result[0].text == "Summary"
assert result[1].type == "text"
parsed = json.loads(result[1].text)
assert parsed == {"data": [1, 2, 3]}
def test_parse_tool_result_from_mcp_structured_content_none():
"""Test that None structuredContent does not affect results."""
mcp_result = types.CallToolResult(
content=[types.TextContent(type="text", text="Hello")],
structuredContent=None,
)
result = _HELPER_MCP_TOOL._parse_tool_result_from_mcp(mcp_result)
assert isinstance(result, list)
assert len(result) == 1
assert result[0].type == "text"
assert result[0].text == "Hello"
def test_parse_tool_result_from_mcp_structured_content_non_serializable():
"""Test that non-JSON-serializable values in structuredContent degrade gracefully."""
mcp_result = types.CallToolResult(
content=[],
structuredContent={"data": b"raw bytes", "count": 42},
)
result = _HELPER_MCP_TOOL._parse_tool_result_from_mcp(mcp_result)
assert isinstance(result, list)
assert len(result) == 1
assert result[0].type == "text"
parsed = json.loads(result[0].text)
assert parsed["count"] == 42
# bytes should be converted to string representation via default=str
assert "raw bytes" in parsed["data"]
def test_mcp_content_types_to_ai_content_text():
"""Test conversion of MCP text content to AI content."""
mcp_content = types.TextContent(type="text", text="Sample text")
@@ -1467,6 +1530,7 @@ def test_mcp_tool_approval_mode_returns_none_for_unmatched_names() -> None:
3,
["tool_one", "tool_two", "tool_three"],
), # None means all tools are allowed
([], 0, []), # Empty list means no tools are allowed
(["tool_one"], 1, ["tool_one"]), # Only tool_one is allowed
(
["tool_one", "tool_three"],
@@ -1813,6 +1877,18 @@ async def test_mcp_tool_message_handler_cancel_and_replace():
assert len(tool._pending_reload_tasks) == 0
def _approve(_params: object) -> bool:
"""Approving sampling gate used by tests that exercise forwarding behavior."""
return True
def _make_sampling_response(text: str = "response", model: str = "test-model") -> Mock:
mock_response = Mock()
mock_response.messages = [Message(role="assistant", contents=[Content.from_text(text)])]
mock_response.model = model
return mock_response
async def test_mcp_tool_sampling_callback_no_client():
"""Test sampling callback error path when no chat client is available."""
tool = MCPStdioTool(name="test_tool", command="python")
@@ -1828,9 +1904,190 @@ async def test_mcp_tool_sampling_callback_no_client():
assert "No chat client available" in result.message
async def test_mcp_tool_sampling_callback_denies_by_default():
"""Sampling is denied when no approval callback is configured (safe default)."""
tool = MCPStdioTool(name="test_tool", command="python")
mock_chat_client = AsyncMock()
tool.client = mock_chat_client
params = Mock()
params.messages = []
params.maxTokens = 128
result = await tool.sampling_callback(Mock(), params)
assert isinstance(result, types.ErrorData)
assert result.code == types.INVALID_REQUEST
assert "denied" in result.message
assert "sampling_approval_callback" in result.message
mock_chat_client.get_response.assert_not_called()
async def test_mcp_tool_sampling_callback_denied_by_callback():
"""Sampling is denied when the approval callback returns a falsy value."""
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=lambda params: False)
mock_chat_client = AsyncMock()
tool.client = mock_chat_client
params = Mock()
params.messages = []
params.maxTokens = 128
result = await tool.sampling_callback(Mock(), params)
assert isinstance(result, types.ErrorData)
assert result.code == types.INVALID_REQUEST
assert "denied by the 'sampling_approval_callback'" in result.message
mock_chat_client.get_response.assert_not_called()
async def test_mcp_tool_sampling_callback_callback_exception_denies():
"""An approval callback that raises results in denial, not an LLM call."""
def boom(_params: object) -> bool:
raise RuntimeError("approval error")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=boom)
mock_chat_client = AsyncMock()
tool.client = mock_chat_client
params = Mock()
params.messages = []
params.maxTokens = 128
result = await tool.sampling_callback(Mock(), params)
assert isinstance(result, types.ErrorData)
assert result.code == types.INVALID_REQUEST
mock_chat_client.get_response.assert_not_called()
async def test_mcp_tool_sampling_callback_async_approval():
"""An async approval callback that approves allows the request through."""
async def approve(_params: object) -> bool:
return True
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=approve)
mock_chat_client = AsyncMock()
mock_chat_client.get_response.return_value = _make_sampling_response("ok")
tool.client = mock_chat_client
params = Mock()
params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
params.temperature = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
params.toolChoice = None
result = await tool.sampling_callback(Mock(), params)
assert isinstance(result, types.CreateMessageResult)
assert result.content.text == "ok"
mock_chat_client.get_response.assert_awaited_once()
async def test_mcp_tool_sampling_callback_clamps_max_tokens():
"""An approved request's maxTokens is clamped to sampling_max_tokens."""
tool = MCPStdioTool(
name="test_tool",
command="python",
sampling_approval_callback=_approve,
sampling_max_tokens=512,
)
mock_chat_client = AsyncMock()
mock_chat_client.get_response.return_value = _make_sampling_response()
tool.client = mock_chat_client
params = Mock()
params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
params.temperature = None
params.maxTokens = 1_000_000
params.stopSequences = None
params.systemPrompt = None
params.tools = None
params.toolChoice = None
result = await tool.sampling_callback(Mock(), params)
assert isinstance(result, types.CreateMessageResult)
options = mock_chat_client.get_response.call_args.kwargs.get("options") or {}
assert options["max_tokens"] == 512
async def test_mcp_tool_sampling_callback_does_not_clamp_under_cap():
"""A request below the cap keeps its requested maxTokens."""
tool = MCPStdioTool(
name="test_tool",
command="python",
sampling_approval_callback=_approve,
sampling_max_tokens=512,
)
mock_chat_client = AsyncMock()
mock_chat_client.get_response.return_value = _make_sampling_response()
tool.client = mock_chat_client
params = Mock()
params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
params.temperature = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
params.toolChoice = None
result = await tool.sampling_callback(Mock(), params)
assert isinstance(result, types.CreateMessageResult)
options = mock_chat_client.get_response.call_args.kwargs.get("options") or {}
assert options["max_tokens"] == 100
async def test_mcp_tool_sampling_callback_rate_limited():
"""Sampling requests beyond sampling_max_requests are rejected per session."""
tool = MCPStdioTool(
name="test_tool",
command="python",
sampling_approval_callback=_approve,
sampling_max_requests=2,
)
mock_chat_client = AsyncMock()
mock_chat_client.get_response.return_value = _make_sampling_response()
tool.client = mock_chat_client
def make_params() -> Mock:
params = Mock()
params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
params.temperature = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
params.toolChoice = None
return params
first = await tool.sampling_callback(Mock(), make_params())
second = await tool.sampling_callback(Mock(), make_params())
third = await tool.sampling_callback(Mock(), make_params())
assert isinstance(first, types.CreateMessageResult)
assert isinstance(second, types.CreateMessageResult)
assert isinstance(third, types.ErrorData)
assert third.code == types.INVALID_REQUEST
assert "rate limit" in third.message.lower()
assert mock_chat_client.get_response.await_count == 2
# The counter resets on a session reset.
tool._reset_session_state()
fourth = await tool.sampling_callback(Mock(), make_params())
assert isinstance(fourth, types.CreateMessageResult)
async def test_mcp_tool_sampling_callback_chat_client_exception():
"""Test sampling callback when chat client raises exception."""
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
# Mock chat client that raises exception
mock_chat_client = AsyncMock()
@@ -1846,7 +2103,7 @@ async def test_mcp_tool_sampling_callback_chat_client_exception():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
@@ -1863,7 +2120,7 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
"""Test sampling callback when response has no valid content types."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
# Mock chat client with response containing only invalid content types
mock_chat_client = AsyncMock()
@@ -1892,7 +2149,7 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
@@ -1905,18 +2162,18 @@ async def test_mcp_tool_sampling_callback_no_valid_content():
assert "Failed to get right content types from the response." in result.message
mock_chat_client.get_response.assert_awaited_once()
_, kwargs = mock_chat_client.get_response.await_args
assert kwargs["options"] == {"max_tokens": None}
assert kwargs["options"] == {"max_tokens": 100}
async def test_mcp_tool_sampling_callback_no_response_and_successful_message_creation():
"""Test sampling callback when the chat client returns no response and then valid content."""
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
tool.client = AsyncMock()
params = Mock()
params.messages = [types.PromptMessage(role="user", content=types.TextContent(type="text", text="Hi"))]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
@@ -1955,7 +2212,7 @@ async def test_mcp_tool_sampling_callback_forwards_system_prompt():
"""Test sampling callback passes systemPrompt as instructions in options."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -1972,7 +2229,7 @@ async def test_mcp_tool_sampling_callback_forwards_system_prompt():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = "You are a helpful assistant"
params.tools = None
@@ -1990,7 +2247,7 @@ async def test_mcp_tool_sampling_callback_forwards_tools():
"""Test sampling callback converts MCP tools to FunctionTools and passes them in options."""
from agent_framework import FunctionTool, Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -2013,7 +2270,7 @@ async def test_mcp_tool_sampling_callback_forwards_tools():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = [mcp_tool]
@@ -2036,7 +2293,7 @@ async def test_mcp_tool_sampling_callback_forwards_tool_choice():
"""Test sampling callback passes toolChoice mode in options."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -2053,7 +2310,7 @@ async def test_mcp_tool_sampling_callback_forwards_tool_choice():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = None
@@ -2071,7 +2328,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_system_prompt():
"""Test sampling callback forwards empty string systemPrompt as instructions."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -2088,7 +2345,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_system_prompt():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = ""
params.tools = None
@@ -2106,7 +2363,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_tools_list():
"""Test sampling callback forwards empty tools list in options."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -2123,7 +2380,7 @@ async def test_mcp_tool_sampling_callback_forwards_empty_tools_list():
mock_message.content.text = "Test question"
params.messages = [mock_message]
params.temperature = None
params.maxTokens = None
params.maxTokens = 100
params.stopSequences = None
params.systemPrompt = None
params.tools = []
@@ -2141,7 +2398,7 @@ async def test_mcp_tool_sampling_callback_forwards_generation_params_in_options(
"""Test sampling callback passes temperature, max_tokens, and stop in options."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -2182,7 +2439,7 @@ async def test_mcp_tool_sampling_callback_omits_temperature_when_none():
"""Test sampling callback does not set temperature in options when it is None."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
@@ -2219,7 +2476,7 @@ async def test_mcp_tool_sampling_callback_always_passes_max_tokens():
"""Test sampling callback always sets max_tokens in options since maxTokens is a required int field."""
from agent_framework import Message
tool = MCPStdioTool(name="test_tool", command="python")
tool = MCPStdioTool(name="test_tool", command="python", sampling_approval_callback=_approve)
mock_chat_client = AsyncMock()
mock_response = Mock()
python/packages/core/tests/core/test_mcp_observability.py
+1
View File
@@ -76,6 +76,7 @@ def _make_call_tool_result(text: str = "result", is_error: bool = False) -> Mock
result = Mock()
result.isError = is_error
result.content = [types.TextContent(type="text", text=text)]
result.structuredContent = None
return result
python/packages/foundry/pyproject.toml
+3 -3
View File
@@ -4,7 +4,7 @@ description = "Microsoft Foundry integrations for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.8.0"
version = "1.8.1"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,8 +23,8 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.8.0,<2",
"agent-framework-openai>=1.8.0,<2",
"agent-framework-core>=1.8.1,<2",
"agent-framework-openai>=1.8.1,<2",
"azure-ai-inference>=1.0.0b9,<1.0.0b10",
"azure-ai-projects>=2.2.0,<3.0",
]
python/packages/foundry_hosting/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "Foundry Hosting integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.0.0a260604"
version = "1.0.0a260609"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,7 +23,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.8.0,<2",
"agent-framework-core>=1.8.1,<2",
"azure-ai-agentserver-core>=2.0.0b3,<3",
"azure-ai-agentserver-responses>=1.0.0b7,<2",
"azure-ai-agentserver-invocations>=1.0.0b3,<2",
python/packages/gemini/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "Google Gemini integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.0.0a260521"
version = "1.0.0a260609"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -24,7 +24,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.6.0,<2.0",
"agent-framework-core>=1.8.1,<2.0",
"google-genai>=1.65.0,<2.0.0",
]
python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py
+4
View File
@@ -978,6 +978,10 @@ async def test_sandbox_code_failure_returns_nonzero_exit(restored_sandbox) -> No
@skip_if_hyperlight_integration_tests_disabled
@pytest.mark.skipif(
sys.platform == "win32" and sys.version_info < (3, 11),
reason="Hyperlight sandbox snapshot/restore crashes on Windows Python 3.10.",
)
async def test_sandbox_snapshot_restore_keeps_sandbox_functional(restored_sandbox) -> None:
"""Verify snapshot/restore cycle leaves the sandbox in a working state."""
# Mutate the sandbox
python/packages/mem0/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "Mem0 integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.0.0b260521"
version = "1.0.0b260609"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,7 +23,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.6.0,<2",
"agent-framework-core>=1.8.1,<2",
"mem0ai>=1.0.0,<2",
]
python/packages/openai/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "OpenAI integrations for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.8.0"
version = "1.8.1"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,7 +23,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core>=1.8.0,<2",
"agent-framework-core>=1.8.1,<2",
"openai>=1.99.0,<3",
]
python/packages/purview/README.md
+1
View File
@@ -320,4 +320,5 @@ except (PurviewAuthenticationError, PurviewRateLimitError, PurviewRequestError,
- **Streaming Responses**: Post-response policy evaluation presently applies only to non-streaming chat responses.
- **Error Handling**: Use `ignore_exceptions` and `ignore_payment_required` settings for graceful degradation. When enabled, errors are logged but don't fail the request.
- **Caching**: Protection scopes responses and 402 errors are cached by default with a 4-hour TTL. Cache is automatically invalidated when protection scope state changes.
- **Cold-cache parallelization**: On a `ProtectionScopes` cache miss, scopes are refreshed in the background while `ProcessContent` runs in the foreground.
- **Background Processing**: Content Activities and offline Process Content requests are handled asynchronously using background tasks to avoid blocking the main execution flow.
python/packages/purview/agent_framework_purview/_processor.py
+58 -27
View File
@@ -231,18 +231,19 @@ class ScopedContentProcessor:
cached_ps_resp = await self._cache.get(cache_key)
if cached_ps_resp is not None and isinstance(cached_ps_resp, ProtectionScopesResponse):
ps_resp = cached_ps_resp
else:
ttl = self._settings.get("cache_ttl_seconds")
ttl_seconds = ttl if ttl is not None else 14400
try:
ps_resp = await self._client.get_protection_scopes(ps_req)
await self._cache.set(cache_key, ps_resp, ttl_seconds=ttl_seconds)
except PurviewPaymentRequiredError as ex:
# Cache the exception at tenant level so all subsequent requests for this tenant fail fast
await self._cache.set(tenant_payment_cache_key, ex, ttl_seconds=ttl_seconds)
raise
return await self._process_with_cached_scopes(pc_request, cached_ps_resp, cache_key)
task = asyncio.create_task(self._refresh_protection_scopes_background(ps_req, cache_key, pc_request))
self._background_tasks.add(task)
task.add_done_callback(self._background_tasks.discard)
return await self._call_process_content(pc_request, cache_key, dlp_actions=[])
async def _process_with_cached_scopes(
self,
pc_request: ProcessContentRequest,
ps_resp: ProtectionScopesResponse,
cache_key: str,
) -> ProcessContentResponse:
if ps_resp.scope_identifier:
pc_request.scope_identifier = ps_resp.scope_identifier
@@ -259,13 +260,7 @@ class ScopedContentProcessor:
task.add_done_callback(self._background_tasks.discard)
return ProcessContentResponse(id="204", correlation_id=pc_request.correlation_id)
pc_resp = await self._client.process_content(pc_request)
if pc_request.scope_identifier and pc_resp.protection_scope_state == ProtectionScopeState.MODIFIED:
await self._cache.remove(cache_key)
pc_resp.policy_actions = self._combine_policy_actions(pc_resp.policy_actions, dlp_actions)
return pc_resp
return await self._call_process_content(pc_request, cache_key, dlp_actions=dlp_actions)
# No applicable scopes - send content activities in background
ca_req = ContentActivitiesRequest(
@@ -281,12 +276,52 @@ class ScopedContentProcessor:
# Respond with HttpStatusCode 204(No Content)
return ProcessContentResponse(id="204", correlation_id=pc_request.correlation_id)
async def _call_process_content(
self,
pc_request: ProcessContentRequest,
cache_key: str,
dlp_actions: list[DlpActionInfo],
) -> ProcessContentResponse:
pc_resp = await self._client.process_content(pc_request)
if pc_request.scope_identifier and pc_resp.protection_scope_state == ProtectionScopeState.MODIFIED:
await self._cache.remove(cache_key)
if dlp_actions:
pc_resp.policy_actions = self._combine_policy_actions(pc_resp.policy_actions, dlp_actions)
return pc_resp
async def _refresh_protection_scopes_background(
self, ps_req: ProtectionScopesRequest, cache_key: str, pc_request: ProcessContentRequest
) -> None:
"""Fetch protection scopes and warm the cache without blocking the foreground call."""
ttl = self._settings.get("cache_ttl_seconds")
ttl_seconds = ttl if ttl is not None else 14400
try:
ps_resp = await self._client.get_protection_scopes(ps_req)
await self._cache.set(cache_key, ps_resp, ttl_seconds=ttl_seconds)
should_process, _, _ = self._check_applicable_scopes(pc_request, ps_resp)
if not should_process:
ca_req = ContentActivitiesRequest(
user_id=pc_request.user_id,
tenant_id=pc_request.tenant_id,
content_to_process=pc_request.content_to_process,
correlation_id=pc_request.correlation_id,
)
await self._send_content_activities_background(ca_req)
except PurviewPaymentRequiredError as ex:
tenant_payment_cache_key = f"purview:payment_required:{ps_req.tenant_id}"
await self._cache.set(tenant_payment_cache_key, ex, ttl_seconds=ttl_seconds)
logger.warning("Background protection scopes refresh failed with payment required: %s", ex)
except Exception as ex:
logger.warning("Background protection scopes refresh failed: %s", ex)
async def _process_content_background(self, pc_request: ProcessContentRequest, cache_key: str) -> None:
"""Process content in background for offline execution mode."""
try:
pc_resp = await self._client.process_content(pc_request)
# If protection scope state is modified, make another PC request and invalidate cache
# If protection scopes changed, invalidate cache and retry once.
if pc_request.scope_identifier and pc_resp.protection_scope_state == ProtectionScopeState.MODIFIED:
await self._cache.remove(cache_key)
await self._client.process_content(pc_request)
@@ -306,14 +341,10 @@ class ScopedContentProcessor:
def _combine_policy_actions(
existing: list[DlpActionInfo] | None, new_actions: list[DlpActionInfo]
) -> list[DlpActionInfo]:
by_key: dict[str, DlpActionInfo] = {}
for a in existing or []:
if a.action:
by_key[a.action] = a
for a in new_actions:
if a.action:
by_key[a.action] = a
return list(by_key.values())
combined: dict[tuple[DlpAction | None, RestrictionAction | None], DlpActionInfo] = {}
for action_info in (existing or []) + new_actions:
combined.setdefault((action_info.action, action_info.restriction_action), action_info)
return list(combined.values())
@staticmethod
def _check_applicable_scopes(
python/packages/purview/tests/purview/test_processor.py
+153 -37
View File
@@ -2,6 +2,7 @@
"""Tests for Purview processor."""
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -217,10 +218,38 @@ class TestScopedContentProcessor:
assert action1 in combined
assert action2 in combined
async def test_combine_policy_actions_preserves_restriction_only_actions(
self, processor: ScopedContentProcessor
) -> None:
"""Test _combine_policy_actions keeps actions that only set restrictionAction."""
existing_action = DlpActionInfo(action=DlpAction.OTHER, restrictionAction=RestrictionAction.OTHER)
restriction_only_action = DlpActionInfo(restriction_action=RestrictionAction.BLOCK)
combined = processor._combine_policy_actions([existing_action], [restriction_only_action])
assert combined == [existing_action, restriction_only_action]
async def test_combine_policy_actions_deduplicates_by_action_and_restriction(
self, processor: ScopedContentProcessor
) -> None:
"""Test _combine_policy_actions removes exact duplicate actions."""
block_action = DlpActionInfo(action=DlpAction.BLOCK_ACCESS, restriction_action=RestrictionAction.BLOCK)
duplicate_block_action = DlpActionInfo(
action=DlpAction.BLOCK_ACCESS, restriction_action=RestrictionAction.BLOCK
)
restriction_only_action = DlpActionInfo(restriction_action=RestrictionAction.BLOCK)
combined = processor._combine_policy_actions(
[block_action],
[duplicate_block_action, restriction_only_action],
)
assert combined == [block_action, restriction_only_action]
async def test_process_with_scopes_calls_client_methods(
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test _process_with_scopes calls get_protection_scopes when scopes response is empty."""
"""Test _process_with_scopes calls process_content immediately and warms scopes in background on cache miss."""
from agent_framework_purview._models import (
ContentActivitiesResponse,
ProtectionScopesResponse,
@@ -236,38 +265,91 @@ class TestScopedContentProcessor:
response = await processor._process_with_scopes(request)
mock_client.get_protection_scopes.assert_called_once()
# When no scopes apply, process_content is not called (activities are sent in background)
mock_client.process_content.assert_not_called()
# The response should have id=204 (No Content) when no scopes apply
assert response.id == "204"
# On cache miss, ProcessContent runs in the foreground and the response is returned.
assert response.id == "response-123"
mock_client.process_content.assert_called_once()
async def test_process_with_scopes_ignores_unexpected_cached_value_type(
# Protection scopes are refreshed in a background task.
await asyncio.gather(*list(processor._background_tasks))
mock_client.get_protection_scopes.assert_called_once()
mock_client.send_content_activities.assert_called_once()
async def test_process_with_scopes_preserves_restriction_only_policy_actions(
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test that a corrupted cache entry does not crash processing."""
"""Test cold-cache ProcessContent actions are not dropped when they only contain restrictionAction."""
from agent_framework_purview._models import ProtectionScopesResponse
request = process_content_request_factory()
restriction_only_action = DlpActionInfo(restriction_action=RestrictionAction.BLOCK)
mock_client.get_protection_scopes = AsyncMock(return_value=ProtectionScopesResponse(**{"value": []}))
mock_client.process_content = AsyncMock(
return_value=ProcessContentResponse(
id="response-123",
protection_scope_state="notModified",
policy_actions=[restriction_only_action],
)
)
response = await processor._process_with_scopes(request)
assert response.policy_actions == [restriction_only_action]
await asyncio.gather(*list(processor._background_tasks))
async def test_process_with_cached_scopes_preserves_restriction_only_policy_actions(
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test cached ProtectionScopes actions are not dropped when they only contain restrictionAction."""
from agent_framework_purview._models import (
ExecutionMode,
PolicyLocation,
PolicyScope,
ProcessContentResponse,
ProtectionScopeActivities,
ProtectionScopesResponse,
)
request = process_content_request_factory()
restriction_only_action = DlpActionInfo(restriction_action=RestrictionAction.BLOCK)
process_content_action = DlpActionInfo(action=DlpAction.OTHER, restriction_action=RestrictionAction.OTHER)
scope_location = PolicyLocation(
data_type="microsoft.graph.policyLocationApplication",
value="app-id",
)
scope = PolicyScope(
activities=ProtectionScopeActivities.UPLOAD_TEXT,
locations=[scope_location],
policy_actions=[restriction_only_action],
execution_mode=ExecutionMode.EVALUATE_INLINE,
)
# Return a valid, inline scope so we stay on the normal (non-background) path.
scope_location = PolicyLocation(**{
"@odata.type": "microsoft.graph.policyLocationApplication",
"value": "app-id",
})
scope = PolicyScope(**{
"activities": ProtectionScopeActivities.UPLOAD_TEXT,
"locations": [scope_location],
"execution_mode": ExecutionMode.EVALUATE_INLINE,
})
mock_client.get_protection_scopes = AsyncMock(return_value=ProtectionScopesResponse(**{"value": [scope]}))
processor._cache.get = AsyncMock(
side_effect=[
None,
ProtectionScopesResponse(scope_identifier="scope-123", scopes=[scope]),
]
) # type: ignore[method-assign]
mock_client.process_content = AsyncMock(
return_value=ProcessContentResponse(
id="response-123",
protection_scope_state="notModified",
policy_actions=[process_content_action],
)
)
response = await processor._process_with_scopes(request)
assert response.policy_actions == [process_content_action, restriction_only_action]
async def test_process_with_scopes_ignores_unexpected_cached_value_type(
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test that a corrupted cache entry does not crash processing."""
from agent_framework_purview._models import ProtectionScopesResponse
request = process_content_request_factory()
mock_client.get_protection_scopes = AsyncMock(return_value=ProtectionScopesResponse(**{"value": []}))
mock_client.process_content = AsyncMock(
return_value=ProcessContentResponse(**{"id": "ok", "protectionScopeState": "notModified"})
)
@@ -279,8 +361,9 @@ class TestScopedContentProcessor:
response = await processor._process_with_scopes(request)
assert response.id == "ok"
mock_client.get_protection_scopes.assert_called_once()
mock_client.process_content.assert_called_once()
await asyncio.gather(*list(processor._background_tasks))
mock_client.get_protection_scopes.assert_called_once()
async def test_process_with_scopes_uses_tenant_payment_exception_cache(
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
@@ -301,8 +384,6 @@ class TestScopedContentProcessor:
self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test offline background processing invalidates cache and retries when scope state changes."""
from agent_framework_purview._models import ProcessContentResponse
request = process_content_request_factory()
request.scope_identifier = "etag-1"
@@ -319,6 +400,36 @@ class TestScopedContentProcessor:
processor._cache.remove.assert_called_once_with("purview:protection_scopes:abc")
assert mock_client.process_content.call_count == 2
async def test_background_scope_refresh_caches_payment_required(
self, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""402 raised during background scope refresh is cached at the tenant level."""
from agent_framework_purview._cache import InMemoryCacheProvider
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
settings = PurviewSettings(
app_name="Test App",
tenant_id="12345678-1234-1234-1234-123456789012",
purview_app_location=PurviewAppLocation(
location_type=PurviewLocationType.APPLICATION, location_value="app-id"
),
)
cache = InMemoryCacheProvider()
processor = ScopedContentProcessor(mock_client, settings, cache_provider=cache)
mock_client.get_protection_scopes = AsyncMock(side_effect=PurviewPaymentRequiredError("nope"))
mock_client.process_content = AsyncMock(
return_value=ProcessContentResponse(**{"id": "pc-1", "protectionScopeState": "notModified"})
)
request = process_content_request_factory()
await processor._process_with_scopes(request)
await asyncio.gather(*list(processor._background_tasks))
cached = await cache.get(f"purview:payment_required:{request.tenant_id}")
assert isinstance(cached, PurviewPaymentRequiredError)
async def test_map_messages_with_user_id_in_additional_properties(self, mock_client: AsyncMock) -> None:
"""Test user_id extraction from message additional_properties."""
settings = PurviewSettings(
@@ -387,6 +498,8 @@ class TestScopedContentProcessor:
self, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test that response is returned when scopes don't apply (activities sent in background)."""
from agent_framework_purview._models import ProtectionScopesResponse
settings = PurviewSettings(
app_name="Test App",
tenant_id="12345678-1234-1234-1234-123456789012",
@@ -398,10 +511,8 @@ class TestScopedContentProcessor:
pc_request = process_content_request_factory()
# Mock get_protection_scopes to return no applicable scopes
mock_ps_response = MagicMock()
mock_ps_response.scopes = []
mock_client.get_protection_scopes.return_value = mock_ps_response
mock_ps_response = ProtectionScopesResponse(scopes=[])
processor._cache.get = AsyncMock(side_effect=[None, mock_ps_response]) # type: ignore[method-assign]
# Mock send_content_activities to return success (called in background)
mock_ca_response = MagicMock()
@@ -410,8 +521,10 @@ class TestScopedContentProcessor:
response = await processor._process_with_scopes(pc_request)
mock_client.get_protection_scopes.assert_called_once()
mock_client.get_protection_scopes.assert_not_called()
mock_client.process_content.assert_not_called()
await asyncio.gather(*list(processor._background_tasks))
mock_client.send_content_activities.assert_called_once()
# Response should have id=204 when no scopes apply
assert response.id == "204"
@@ -419,6 +532,8 @@ class TestScopedContentProcessor:
self, mock_client: AsyncMock, process_content_request_factory
) -> None:
"""Test that errors in background activities don't affect the response."""
from agent_framework_purview._models import ProtectionScopesResponse
settings = PurviewSettings(
app_name="Test App",
tenant_id="12345678-1234-1234-1234-123456789012",
@@ -430,10 +545,8 @@ class TestScopedContentProcessor:
pc_request = process_content_request_factory()
# Mock get_protection_scopes to return no applicable scopes
mock_ps_response = MagicMock()
mock_ps_response.scopes = []
mock_client.get_protection_scopes.return_value = mock_ps_response
mock_ps_response = ProtectionScopesResponse(scopes=[])
processor._cache.get = AsyncMock(side_effect=[None, mock_ps_response]) # type: ignore[method-assign]
# Mock send_content_activities to return error (called in background task)
mock_ca_response = MagicMock()
@@ -445,6 +558,8 @@ class TestScopedContentProcessor:
# Since activities are sent in background, errors don't affect the response
# Response should have id=204 when no scopes apply
assert response.id == "204"
await asyncio.gather(*list(processor._background_tasks))
mock_client.send_content_activities.assert_called_once()
class TestUserIdResolution:
@@ -656,10 +771,12 @@ class TestScopedContentProcessorCaching:
mock_client.get_protection_scopes.return_value = ProtectionScopesResponse(
scope_identifier="scope-123", scopes=[]
)
mock_client.process_content.return_value = ProcessContentResponse(id="ok", protection_scope_state="notModified")
messages = [Message(role="user", contents=["Test"])]
await processor.process_messages(messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012")
await asyncio.gather(*list(processor._background_tasks))
mock_client.get_protection_scopes.assert_called_once()
@@ -670,7 +787,7 @@ class TestScopedContentProcessorCaching:
async def test_payment_required_exception_cached_at_tenant_level(
self, mock_client: AsyncMock, settings: PurviewSettings
) -> None:
"""Test that 402 payment required exceptions are cached at tenant level."""
"""Test that background scope 402 returns once, then throws from the tenant-level cache."""
from agent_framework_purview._cache import InMemoryCacheProvider
from agent_framework_purview._exceptions import PurviewPaymentRequiredError
@@ -678,13 +795,12 @@ class TestScopedContentProcessorCaching:
processor = ScopedContentProcessor(mock_client, settings, cache_provider=cache_provider)
mock_client.get_protection_scopes.side_effect = PurviewPaymentRequiredError("Payment required")
mock_client.process_content.return_value = ProcessContentResponse(id="ok", protection_scope_state="notModified")
messages = [Message(role="user", contents=["Test"])]
with pytest.raises(PurviewPaymentRequiredError):
await processor.process_messages(
messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012"
)
await processor.process_messages(messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012")
await asyncio.gather(*list(processor._background_tasks))
mock_client.get_protection_scopes.assert_called_once()
python/pyproject.toml
+2 -2
View File
@@ -4,7 +4,7 @@ description = "Microsoft Agent Framework for building AI Agents with Python. Thi
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "1.8.0"
version = "1.8.1"
license-files = ["LICENSE"]
urls.homepage = "https://aka.ms/agent-framework"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -23,7 +23,7 @@ classifiers = [
"Typing :: Typed",
]
dependencies = [
"agent-framework-core[all]==1.8.0",
"agent-framework-core[all]==1.8.1",
]
[dependency-groups]
python/samples/02-agents/harness/README.md
+35 -2
View File
@@ -17,6 +17,7 @@ from a chat client.
| AgentModeProvider | Plan/execute mode tracking |
| MemoryContextProvider | File-based durable memory (when `memory_store` provided) |
| SkillsProvider | File-based skill discovery and progressive loading |
| Shell tool | Shell command execution + environment probing (when `shell_executor` provided) |
| OpenTelemetry | Built-in observability |
Each feature can be disabled or customized via keyword arguments.
@@ -45,13 +46,23 @@ python samples/02-agents/harness/harness_research.py
### Minimal Setup
`create_harness_agent` requires only a chat client and token budget parameters:
`create_harness_agent` requires only a chat client:
```python
from agent_framework import create_harness_agent
from agent_framework.foundry import FoundryChatClient
from azure.identity import AzureCliCredential
agent = create_harness_agent(
client=FoundryChatClient(credential=AzureCliCredential()),
)
```
### With Compaction
Provide token budget parameters to enable automatic context-window compaction:
```python
agent = create_harness_agent(
client=FoundryChatClient(credential=AzureCliCredential()),
max_context_window_tokens=128_000,
@@ -59,7 +70,7 @@ agent = create_harness_agent(
)
```
### Customization
### Further Customization
Disable or customize any feature:
@@ -81,3 +92,25 @@ agent = create_harness_agent(
The `AgentModeProvider` enables a two-phase workflow:
1. **Plan mode** — Interactive: the agent asks questions, creates todos, gets approval
2. **Execute mode** — Autonomous: the agent works through todos independently
### Shell Tool
Pass a shell executor (e.g. `LocalShellTool` from `agent-framework-tools`) to enable shell
command execution plus automatic environment probing via a `ShellEnvironmentProvider`. The
tool is only wired when the chat client supports shell tools; otherwise a warning is logged
and the shell tool/provider are skipped. The caller owns the executor's lifecycle.
```python
from agent_framework_tools.shell import LocalShellTool, ShellEnvironmentProviderOptions
async with LocalShellTool(acknowledge_unsafe=True) as shell:
agent = create_harness_agent(
client=client,
max_context_window_tokens=128_000,
max_output_tokens=16_384,
shell_executor=shell,
# Optional: customize environment probing.
shell_environment_provider_options=ShellEnvironmentProviderOptions(probe_tools=("git", "python")),
)
```
python/samples/02-agents/harness/console/agent_runner.py
+1 -3
View File
@@ -313,9 +313,7 @@ class HarnessAgentRunner:
"""
actions: list[FollowUpAction] = []
for observer in self._observers:
observer_actions = await observer.on_stream_complete(
self._ux, self._agent, session
)
observer_actions = await observer.on_stream_complete(self._ux, self._agent, session)
if observer_actions:
actions.extend(observer_actions)
return actions
python/samples/02-agents/harness/console/app.py
+2 -8
View File
@@ -182,18 +182,12 @@ class HarnessApp(App[None]):
if command_handlers is None:
from .commands import build_default_command_handlers
self._command_handlers = build_default_command_handlers(
agent, mode_colors=mode_colors
)
self._command_handlers = build_default_command_handlers(agent, mode_colors=mode_colors)
else:
self._command_handlers = command_handlers
# Compute help text from command handlers
help_parts = [
h.get_help_text()
for h in self._command_handlers
if h.get_help_text() is not None
]
help_parts = [h.get_help_text() for h in self._command_handlers if h.get_help_text() is not None]
help_text = ", ".join(help_parts) if help_parts else None
# State and driver
python/samples/02-agents/harness/console/commands/todo_handler.py
+1 -3
View File
@@ -45,9 +45,7 @@ class TodoCommandHandler(CommandHandler):
ux.append_info_line("TodoProvider is not available.")
return True
todos = await self._todo_provider.store.load_items(
session, source_id=self._todo_provider.source_id
)
todos = await self._todo_provider.store.load_items(session, source_id=self._todo_provider.source_id)
if not todos:
ux.append_info_line("No todos yet.")
python/samples/02-agents/harness/console/components/scroll_panel.py
+1 -1
View File
@@ -72,7 +72,7 @@ class HarnessScrollPanel(RichLog):
# Truncate lines back to where streaming started
if len(self.lines) > self._streaming_line_start:
del self.lines[self._streaming_line_start:]
del self.lines[self._streaming_line_start :]
from textual.geometry import Size
self.virtual_size = Size(self._widest_line_width, len(self.lines))
python/samples/02-agents/harness/console/observers/planning_models.py
+1 -2
View File
@@ -41,8 +41,7 @@ class PlanningQuestion(BaseModel):
choices: list[str] | None = Field(
default=None,
description=(
"For clarifications, this has a list of options that the user can "
"choose from. null for approvals."
"For clarifications, this has a list of options that the user can choose from. null for approvals."
),
)
python/samples/02-agents/mcp/README.md
+1
View File
@@ -14,6 +14,7 @@ The Model Context Protocol (MCP) is an open standard for connecting AI agents to
| **API Key Authentication** | [`mcp_api_key_auth.py`](mcp_api_key_auth.py) | Demonstrates API key authentication with MCP servers using `header_provider`, runtime invocation kwargs, and a command-line API key argument |
| **GitHub Integration with PAT** | [`mcp_github_pat.py`](mcp_github_pat.py) | Demonstrates connecting to GitHub's MCP server using Personal Access Token (PAT) authentication |
| **Long-Running Task** | [`mcp_long_running_task.py`](mcp_long_running_task.py) | Demonstrates transparent SEP-2663 long-running task handling for MCP tools that advertise `taskSupport=required`. Self-spawns a stdio MCP child server |
| **Sampling Approval** | [`mcp_sampling_approval.py`](mcp_sampling_approval.py) | Demonstrates gating server-initiated `sampling/createMessage` requests with a `sampling_approval_callback`, plus the `sampling_max_tokens` and `sampling_max_requests` guardrails. MCP sampling is denied by default |
## Prerequisites
python/samples/02-agents/mcp/mcp_sampling_approval.py
+78
View File
@@ -0,0 +1,78 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import Agent, MCPStreamableHTTPTool
from agent_framework.openai import OpenAIChatClient
from dotenv import load_dotenv
from mcp import types
# Load environment variables from .env file
load_dotenv()
"""
MCP Sampling Approval Example
MCP servers can send the client a ``sampling/createMessage`` request, asking the
client to run an LLM completion on the server's behalf. Because remote MCP
servers are untrusted third parties, forwarding these server-controlled prompts
to your chat client without review is a confused-deputy risk: a malicious server
could exfiltrate context, force tool calls, or burn through your token budget.
For that reason Agent Framework **denies MCP sampling by default**. To allow it,
pass a ``sampling_approval_callback`` to the MCP tool. The callback receives the
raw ``CreateMessageRequestParams`` and returns ``True`` to approve or ``False``
to deny. It may be synchronous or asynchronous, so you can implement a
human-in-the-loop prompt, a policy check, or an audit log.
Two further guardrails apply to approved requests:
- ``sampling_max_tokens`` caps the server-requested ``maxTokens``.
- ``sampling_max_requests`` limits how many sampling requests a single session
may make.
To restore the legacy "always approve" behavior (only do this for servers you
trust), pass ``sampling_approval_callback=lambda params: True``.
"""
async def approve_sampling(params: types.CreateMessageRequestParams) -> bool:
"""Human-in-the-loop approval gate for server-initiated sampling.
Shows the server-supplied system prompt and messages, then asks the user to
approve or deny. Returning ``False`` rejects the request.
"""
print("\n--- MCP server requested a sampling/createMessage ---")
if params.systemPrompt:
print(f"System prompt: {params.systemPrompt}")
for message in params.messages:
text = getattr(message.content, "text", message.content)
print(f"{message.role}: {text}")
answer = await asyncio.to_thread(input, "Approve this sampling request? [y/N]: ")
return answer.strip().lower() in {"y", "yes"}
async def main() -> None:
"""Run an agent against an MCP server with a sampling approval gate."""
async with Agent(
client=OpenAIChatClient(),
name="Agent",
instructions="You are a helpful assistant. Use your MCP tool when answering the user's question.",
tools=MCPStreamableHTTPTool(
name="MCP tool",
description="MCP tool description.",
url="<your mcp server url>",
# Passing ``client`` enables sampling; the approval callback gates it.
client=OpenAIChatClient(),
sampling_approval_callback=approve_sampling,
sampling_max_tokens=2048,
sampling_max_requests=5,
),
) as agent:
query = "Use your MCP tool to help answer this question."
print(f"User: {query}")
result = await agent.run(query)
print(f"Agent: {result.text}")
if __name__ == "__main__":
asyncio.run(main())
python/samples/02-agents/tools/README.md
+1
View File
@@ -22,6 +22,7 @@ injection, and dynamic (progressive) tool exposure.
|------|--------------|
| [`function_tool_with_approval.py`](function_tool_with_approval.py) | Requiring human approval before a tool runs. |
| [`function_tool_with_approval_and_sessions.py`](function_tool_with_approval_and_sessions.py) | Tool approvals combined with sessions. |
| [`tool_approval_middleware.py`](tool_approval_middleware.py) | Session-backed approval coordination, mixed-batch approvals, and "always approve" rules. |
| [`function_invocation_configuration.py`](function_invocation_configuration.py) | Configuring function-invocation settings (e.g. max iterations). |
| [`control_total_tool_executions.py`](control_total_tool_executions.py) | All the ways to cap how many times tools run. |
| [`function_tool_with_max_invocations.py`](function_tool_with_max_invocations.py) | Limiting the number of invocations per tool. |
python/samples/02-agents/tools/tool_approval_middleware.py
+191
View File
@@ -0,0 +1,191 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from typing import Annotated
from agent_framework import (
Agent,
AgentResponse,
AgentSession,
Content,
Message,
ToolApprovalMiddleware,
create_always_approve_tool_response,
create_always_approve_tool_with_arguments_response,
tool,
)
from agent_framework.foundry import FoundryChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
"""
This sample demonstrates how a host application can decide which approval
requests may run now, which must be rejected, and which can be remembered for
future runs.
The model may not request every tool on every run. The important part is the
approval mechanism:
1. Tools that are safe to run immediately use ``approval_mode="never_require"``.
2. Sensitive tools use ``approval_mode="always_require"``.
3. ``ToolApprovalMiddleware`` coordinates approval prompts and standing rules.
4. The host turns user policy into ``function_approval_response`` content:
- approve for this request only;
- reject for this request;
- approve and remember the tool for future requests;
- approve and remember the tool only when called again with the same arguments.
5. Heuristic auto-approval rules can approve low-risk function calls before
the user is prompted.
"""
# Load environment variables from .env file
load_dotenv()
@tool(approval_mode="never_require")
def lookup_ticket(ticket_id: Annotated[str, "Support ticket id, for example T-123"]) -> str:
"""Look up a support ticket. This read-only tool runs without approval."""
return f"Ticket {ticket_id}: customer confirmed the issue can be closed."
@tool(approval_mode="always_require")
def close_ticket(
ticket_id: Annotated[str, "Support ticket id, for example T-123"],
resolution: Annotated[str, "Short resolution text"],
) -> str:
"""Close a support ticket."""
return f"Ticket {ticket_id} closed with resolution: {resolution}"
@tool(approval_mode="always_require")
def notify_customer(
ticket_id: Annotated[str, "Support ticket id, for example T-123"],
message: Annotated[str, "Message to send to the customer"],
) -> str:
"""Notify the customer about a ticket update."""
return f"Customer notified for {ticket_id}: {message}"
@tool(approval_mode="always_require")
def add_internal_note(
ticket_id: Annotated[str, "Support ticket id, for example T-123"],
note: Annotated[str, "Internal note text"],
) -> str:
"""Add an internal note to a support ticket."""
return f"Internal note added to {ticket_id}: {note}"
@tool(approval_mode="always_require")
def delete_attachment(
ticket_id: Annotated[str, "Support ticket id, for example T-123"],
attachment_name: Annotated[str, "Attachment file name"],
) -> str:
"""Delete an attachment from a support ticket."""
return f"Deleted {attachment_name} from ticket {ticket_id}."
def auto_approve_low_risk_notes(function_call: Content) -> bool:
"""Heuristic rule: auto-approve short internal notes for the target ticket."""
if function_call.name != "add_internal_note":
return False
arguments = function_call.parse_arguments() or {}
note = str(arguments.get("note", ""))
return arguments.get("ticket_id") == "T-123" and len(note) <= 120
def approval_response_for_user_policy(request: Content) -> Content:
"""Convert user/host policy into an approval response for one tool request."""
function_call = request.function_call
if function_call is None or function_call.name is None:
return request.to_function_approval_response(approved=False)
tool_name = function_call.name
print(f"Approval requested: {tool_name}({function_call.arguments})")
if tool_name in {"close_ticket"}:
print(f"Decision: approve and remember future {tool_name} calls with these exact arguments")
return create_always_approve_tool_with_arguments_response(request)
if tool_name in {"notify_customer"}:
print(f"Decision: approve and remember all future {tool_name} calls")
return create_always_approve_tool_response(request)
if tool_name in {"delete_attachment"}:
print(f"Decision: reject {tool_name} for this run")
return request.to_function_approval_response(approved=False)
print(f"Decision: reject {tool_name}; no policy allowed it")
return request.to_function_approval_response(approved=False)
async def resolve_approval_requests(agent: Agent, response: AgentResponse, session: AgentSession) -> AgentResponse:
"""Resolve approval prompts until the agent returns a regular answer."""
result = response
while result.user_input_requests:
approval_responses = [approval_response_for_user_policy(request) for request in result.user_input_requests]
result = await agent.run(Message(role="user", contents=approval_responses), session=session)
return result
async def main() -> None:
"""Run the tool approval middleware sample."""
# 1. Create a regular chat client.
client = FoundryChatClient(credential=AzureCliCredential())
# 2. Create an agent with sensitive tools and opt-in ToolApprovalMiddleware.
agent = Agent(
client=client,
name="SupportAgent",
instructions=(
"You are a support agent. Use tools when useful. "
"Look up ticket T-123, close it if the customer confirmed, notify the customer, "
"add a short internal note, and do not delete attachments unless the tool is approved."
),
tools=[lookup_ticket, close_ticket, notify_customer, add_internal_note, delete_attachment],
middleware=[ToolApprovalMiddleware(auto_approval_rules=[auto_approve_low_risk_notes])],
)
session = agent.create_session()
# 3. Ask for work that may trigger a mixed batch of safe and sensitive tool calls.
query = (
"Please process ticket T-123: check the ticket, close it as resolved, "
"notify the customer, add a short internal note, and remove debug.log if it is attached."
)
print(f"User: {query}")
result = await agent.run(query, session=session)
# 4. Convert approval requests into approve/reject/always-approve responses.
result = await resolve_approval_requests(agent, result, session)
print(f"Agent: {result.text}")
# 5. Later runs can use remembered approval rules:
# - notify_customer: all future calls to the tool.
# - close_ticket: only future calls with the same arguments.
# - add_internal_note: low-risk matching calls are auto-approved by the heuristic callback.
follow_up = "Send the customer a short follow-up for ticket T-123."
print(f"\nUser: {follow_up}")
result = await agent.run(follow_up, session=session)
result = await resolve_approval_requests(agent, result, session)
print(f"Agent: {result.text}")
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output:
User: Please process ticket T-123: check the ticket, close it as resolved,
notify the customer, add a short internal note, and remove debug.log if it is attached.
Approval requested: close_ticket({"ticket_id": "T-123", "resolution": "resolved"})
Decision: approve and remember future close_ticket calls with these exact arguments
Approval requested: notify_customer({"ticket_id": "T-123", "message": "Your ticket has been resolved."})
Decision: approve and remember all future notify_customer calls
Approval requested: delete_attachment({"ticket_id": "T-123", "attachment_name": "debug.log"})
Decision: reject delete_attachment for this run
Agent: Ticket T-123 was closed, the customer was notified, and a short internal note was added.
I did not delete debug.log.
User: Send the customer a short follow-up for ticket T-123.
Agent: The customer was sent a short follow-up for ticket T-123.
"""
python/samples/05-end-to-end/purview_agent/README.md
+34 -10
View File
@@ -3,7 +3,7 @@
This getting-started sample shows how to attach Microsoft Purview policy evaluation to an Agent Framework `Agent` using the **middleware** approach.
**What this sample demonstrates:**
1. Configure an Azure OpenAI chat client
1. Configure a Foundry chat client
2. Add Purview policy enforcement middleware (`PurviewPolicyMiddleware`)
3. Add Purview policy enforcement at the chat client level (`PurviewChatPolicyMiddleware`)
4. Implement a custom cache provider for advanced caching scenarios
@@ -17,8 +17,8 @@ This getting-started sample shows how to attach Microsoft Purview policy evaluat
| Variable | Required | Purpose |
|----------|----------|---------|
| `AZURE_OPENAI_ENDPOINT` | Yes | Azure OpenAI endpoint (https://<name>.openai.azure.com) |
| `AZURE_OPENAI_MODEL` | Optional | Model deployment name (defaults inside SDK if omitted) |
| `FOUNDRY_PROJECT_ENDPOINT` | Yes | Azure AI Foundry project endpoint, for example `https://<resource>.services.ai.azure.com/api/projects/<project>` |
| `FOUNDRY_MODEL` | Optional | Model deployment name (defaults to `gpt-4o-mini`) |
| `PURVIEW_CLIENT_APP_ID` | Yes* | Client (application) ID used for Purview authentication |
| `PURVIEW_USE_CERT_AUTH` | Optional (`true`/`false`) | Switch between certificate and interactive auth |
| `PURVIEW_TENANT_ID` | Yes (when cert auth on) | Tenant ID for certificate authentication |
@@ -31,7 +31,8 @@ This getting-started sample shows how to attach Microsoft Purview policy evaluat
Opens a browser on first run to sign in.
```powershell
$env:AZURE_OPENAI_ENDPOINT = "https://your-openai-instance.openai.azure.com"
$env:FOUNDRY_PROJECT_ENDPOINT = "https://<resource>.services.ai.azure.com/api/projects/<project>"
$env:FOUNDRY_MODEL = "gpt-4o-mini"
$env:PURVIEW_CLIENT_APP_ID = "00000000-0000-0000-0000-000000000000"
```
@@ -64,22 +65,27 @@ If interactive auth is used, a browser window will appear the first time.
## 4. How It Works
The sample demonstrates three different scenarios:
The sample demonstrates four integration scenarios. Each scenario runs the same three-message sequence via `run_policy_flow(...)`:
1. **good (cold cache)** - a benign prompt that exercises the cold-cache parallel ProtectionScopes warmup + foreground ProcessContent path.
2. **expected block** - a sensitive prompt containing the Visa test credit card number `4111 1111 1111 1111`. If the tenant has a DLP policy for `Microsoft 365 Copilot and AI apps` targeting the Credit Card sensitive info type with a Block action, this prompt returns the configured `blocked_prompt_message` (default: `Prompt blocked by policy`). If no DLP policy applies, the prompt is allowed (the LLM may still decline on its own, but that is a model-level response, not a Purview block).
3. **good (warm cache)** - a second benign prompt that exercises the warm-cache path. The custom cache provider scenario prints `Cache HIT` for the same protection-scopes key, confirming the cache and middleware state survive a prior block.
### A. Agent Middleware (`run_with_agent_middleware`)
1. Builds an Azure OpenAI chat client (using the environment endpoint / deployment)
1. Builds a Foundry chat client (using the environment project endpoint / deployment)
2. Chooses credential mode (certificate vs interactive)
3. Creates `PurviewPolicyMiddleware` with `PurviewSettings`
4. Injects middleware into the agent at construction
5. Sends two user messages sequentially
6. Prints results (or policy block messages)
5. Runs the three-message `good -> block -> good` orchestration
6. Prints `ALLOWED` or `BLOCKED` per message, plus the model response
7. Uses default caching automatically
### B. Chat Client Middleware (`run_with_chat_middleware`)
1. Creates a chat client with `PurviewChatPolicyMiddleware` attached directly
2. Policy evaluation happens at the chat client level rather than agent level
3. Demonstrates an alternative integration point for Purview policies
4. Uses default caching automatically
4. Runs the same `good -> block -> good` orchestration
5. Uses default caching automatically
### C. Custom Cache Provider (`run_with_custom_cache_provider`)
1. Implements the `CacheProvider` protocol with a custom class (`SimpleDictCacheProvider`)
@@ -88,9 +94,27 @@ The sample demonstrates three different scenarios:
- `async def get(self, key: str) -> Any | None`
- `async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None`
- `async def remove(self, key: str) -> None`
4. Runs the `good -> block -> good` orchestration and prints `Cache MISS`/`Cache HIT` traces alongside policy outcomes, showing the cold-cache warmup populating the cache and warm-cache requests skipping ProtectionScopes.
### D. Default Cache (`run_with_default_cache`)
1. Same as the agent middleware path but with explicit cache TTL and size limits in `PurviewSettings`
2. Uses the default in-memory `CacheProvider`
3. Runs the `good -> block -> good` orchestration
**Policy Behavior:**
Prompt blocks set a system-level message: `Prompt blocked by policy` and terminate the run early. Response blocks rewrite the output to `Response blocked by policy`.
Prompt blocks substitute the configured `blocked_prompt_message` (default `Prompt blocked by policy`) and terminate the agent run early. Response blocks substitute `blocked_response_message`. The LLM is never called for a blocked prompt.
**Seeing a real `BLOCKED` outcome:**
The middle prompt only returns `BLOCKED` if the tenant actually has a Purview DLP policy that matches the request. Specifically, all of the following must be true:
1. The Entra app id used by `PURVIEW_CLIENT_APP_ID` (the same id Agent Framework sends as `policyLocationApplication.value`) is registered as an integrated AI app in Purview (Settings -> AI app and agent locations).
2. A DLP policy in the tenant targets the location `Microsoft 365 Copilot and AI apps`, scoped to that app id (or `All apps`).
3. The policy has a rule with the condition `Content contains -> Sensitive info types -> Credit Card Number` and an action of `Restrict access to Microsoft 365 Copilot and AI apps -> Block`.
4. The policy is `On` (not `Test mode without notifications`).
5. The signed-in user is in the policy's user scope.
6. Required Graph delegated permissions are admin-consented: `ProtectionScopes.Compute.All`, `Content.Process.All`, `ContentActivity.Write`.
If any of those are missing, the credit card prompt is allowed at the Purview layer. The model itself may still decline on its own; that response is a model-level refusal, not a Purview block. The cold/warm cache orchestration is still demonstrated either way - the `Cache MISS -> Cache HIT` trace from the custom cache scenario does not depend on a block firing.
---
python/samples/05-end-to-end/purview_agent/sample_purview_agent.py
+74 -89
View File
@@ -11,8 +11,8 @@ Shows:
Note: Caching is automatic and enabled by default.
Environment variables:
- AZURE_OPENAI_ENDPOINT (required)
- AZURE_OPENAI_MODEL (optional, defaults to gpt-4o-mini)
- FOUNDRY_PROJECT_ENDPOINT (required) - Azure AI Foundry project endpoint URL
- FOUNDRY_MODEL (optional, defaults to gpt-4o-mini)
- PURVIEW_CLIENT_APP_ID (required)
- PURVIEW_USE_CERT_AUTH (optional, set to "true" for certificate auth)
- PURVIEW_TENANT_ID (required if certificate auth)
@@ -45,6 +45,37 @@ load_dotenv()
JOKER_NAME = "Joker"
JOKER_INSTRUCTIONS = "You are good at telling jokes. Keep responses concise."
# Sequential prompts to demonstrate good -> block -> good orchestration.
# The sensitive prompt contains a Visa test credit card number that matches Purview's
# built-in Credit Card sensitive information type. If the tenant has a DLP policy that
# blocks credit card content for Microsoft 365 Copilot and AI apps, the second message
# will be blocked and the third will verify that subsequent calls still flow normally
# after a block.
GOOD_PROMPT_PRIMARY = "Tell me a joke about a pirate."
SENSITIVE_PROMPT = "My corporate credit card is 4111 1111 1111 1111. Please confirm receipt."
GOOD_PROMPT_FOLLOWUP = "Another light joke please."
async def run_policy_flow(
label: str,
agent: Agent,
user_id: str | None,
blocked_text: str,
) -> None:
"""Run a good -> block candidate -> good sequence and report each outcome."""
blocked_marker = blocked_text.lower()
prompts = [
("good (cold cache)", GOOD_PROMPT_PRIMARY),
("expected block", SENSITIVE_PROMPT),
("good (warm cache)", GOOD_PROMPT_FOLLOWUP),
]
for tag, text in prompts:
response: AgentResponse = await agent.run(
Message("user", [text], additional_properties={"user_id": user_id})
)
outcome = "BLOCKED" if blocked_marker in str(response).lower() else "ALLOWED"
print(f"[{label}] {tag}: {outcome}\n{response}\n")
# Custom Cache Provider Implementation
class SimpleDictCacheProvider:
@@ -138,21 +169,17 @@ def build_credential() -> Any:
async def run_with_agent_middleware() -> None:
endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
endpoint = os.environ.get("FOUNDRY_PROJECT_ENDPOINT")
if not endpoint:
print("Skipping run: AZURE_OPENAI_ENDPOINT not set")
print("Skipping run: FOUNDRY_PROJECT_ENDPOINT not set")
return
deployment = os.environ.get("AZURE_OPENAI_MODEL", "gpt-4o-mini")
deployment = os.environ.get("FOUNDRY_MODEL", "gpt-4o-mini")
user_id = os.environ.get("PURVIEW_DEFAULT_USER_ID")
client = FoundryChatClient(model=deployment, endpoint=endpoint, credential=AzureCliCredential())
client = FoundryChatClient(model=deployment, project_endpoint=endpoint, credential=AzureCliCredential())
purview_agent_middleware = PurviewPolicyMiddleware(
build_credential(),
PurviewSettings(
app_name="Agent Framework Sample App",
),
)
settings = PurviewSettings(app_name="Agent Framework Sample App")
purview_agent_middleware = PurviewPolicyMiddleware(build_credential(), settings)
agent = Agent(
client=client,
@@ -162,39 +189,26 @@ async def run_with_agent_middleware() -> None:
)
print("-- Agent MiddlewareTypes Path --")
first: AgentResponse = await agent.run(
Message("user", ["Tell me a joke about a pirate."], additional_properties={"user_id": user_id})
)
print("First response (agent middleware):\n", first)
second: AgentResponse = await agent.run(
Message(
role="user", contents=["That was funny. Tell me another one."], additional_properties={"user_id": user_id}
)
)
print("Second response (agent middleware):\n", second)
blocked_text = settings.get("blocked_prompt_message") or "Prompt blocked by policy"
await run_policy_flow("agent middleware", agent, user_id, blocked_text)
async def run_with_chat_middleware() -> None:
endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
endpoint = os.environ.get("FOUNDRY_PROJECT_ENDPOINT")
if not endpoint:
print("Skipping chat middleware run: AZURE_OPENAI_ENDPOINT not set")
print("Skipping chat middleware run: FOUNDRY_PROJECT_ENDPOINT not set")
return
deployment = os.environ.get("AZURE_OPENAI_MODEL", default="gpt-4o-mini")
deployment = os.environ.get("FOUNDRY_MODEL", default="gpt-4o-mini")
user_id = os.environ.get("PURVIEW_DEFAULT_USER_ID")
settings = PurviewSettings(app_name="Agent Framework Sample App (Chat)")
client = FoundryChatClient(
model=deployment,
endpoint=endpoint,
project_endpoint=endpoint,
credential=AzureCliCredential(),
middleware=[
PurviewChatPolicyMiddleware(
build_credential(),
PurviewSettings(
app_name="Agent Framework Sample App (Chat)",
),
)
PurviewChatPolicyMiddleware(build_credential(), settings)
],
)
@@ -205,43 +219,27 @@ async def run_with_chat_middleware() -> None:
)
print("-- Chat MiddlewareTypes Path --")
first: AgentResponse = await agent.run(
Message(
role="user",
contents=["Give me a short clean joke."],
additional_properties={"user_id": user_id},
)
)
print("First response (chat middleware):\n", first)
second: AgentResponse = await agent.run(
Message(
role="user",
contents=["One more please."],
additional_properties={"user_id": user_id},
)
)
print("Second response (chat middleware):\n", second)
blocked_text = settings.get("blocked_prompt_message") or "Prompt blocked by policy"
await run_policy_flow("chat middleware", agent, user_id, blocked_text)
async def run_with_custom_cache_provider() -> None:
"""Demonstrate implementing and using a custom cache provider."""
endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
endpoint = os.environ.get("FOUNDRY_PROJECT_ENDPOINT")
if not endpoint:
print("Skipping custom cache provider run: AZURE_OPENAI_ENDPOINT not set")
print("Skipping custom cache provider run: FOUNDRY_PROJECT_ENDPOINT not set")
return
deployment = os.environ.get("AZURE_OPENAI_MODEL", "gpt-4o-mini")
deployment = os.environ.get("FOUNDRY_MODEL", "gpt-4o-mini")
user_id = os.environ.get("PURVIEW_DEFAULT_USER_ID")
client = FoundryChatClient(model=deployment, endpoint=endpoint, credential=AzureCliCredential())
client = FoundryChatClient(model=deployment, project_endpoint=endpoint, credential=AzureCliCredential())
custom_cache = SimpleDictCacheProvider()
settings = PurviewSettings(app_name="Agent Framework Sample App (Custom Provider)")
purview_agent_middleware = PurviewPolicyMiddleware(
build_credential(),
PurviewSettings(
app_name="Agent Framework Sample App (Custom Provider)",
),
settings,
cache_provider=custom_cache,
)
@@ -254,38 +252,28 @@ async def run_with_custom_cache_provider() -> None:
print("-- Custom Cache Provider Path --")
print("Using SimpleDictCacheProvider")
blocked_text = settings.get("blocked_prompt_message") or "Prompt blocked by policy"
await run_policy_flow("custom cache", agent, user_id, blocked_text)
first: AgentResponse = await agent.run(
Message(
role="user", contents=["Tell me a joke about a programmer."], additional_properties={"user_id": user_id}
)
)
print("First response (custom provider):\n", first)
second: AgentResponse = await agent.run(
Message("user", ["That's hilarious! One more?"], additional_properties={"user_id": user_id})
)
print("Second response (custom provider):\n", second)
async def run_with_default_cache() -> None:
"""Demonstrate using the default built-in cache."""
endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
endpoint = os.environ.get("FOUNDRY_PROJECT_ENDPOINT")
if not endpoint:
print("Skipping default cache run: AZURE_OPENAI_ENDPOINT not set")
print("Skipping default cache run: FOUNDRY_PROJECT_ENDPOINT not set")
return
deployment = os.environ.get("AZURE_OPENAI_MODEL", "gpt-4o-mini")
deployment = os.environ.get("FOUNDRY_MODEL", "gpt-4o-mini")
user_id = os.environ.get("PURVIEW_DEFAULT_USER_ID")
client = FoundryChatClient(model=deployment, endpoint=endpoint, credential=AzureCliCredential())
client = FoundryChatClient(model=deployment, project_endpoint=endpoint, credential=AzureCliCredential())
# No cache_provider specified - uses default InMemoryCacheProvider
purview_agent_middleware = PurviewPolicyMiddleware(
build_credential(),
PurviewSettings(
app_name="Agent Framework Sample App (Default Cache)",
cache_ttl_seconds=3600,
max_cache_size_bytes=100 * 1024 * 1024, # 100MB
),
settings = PurviewSettings(
app_name="Agent Framework Sample App (Default Cache)",
cache_ttl_seconds=3600,
max_cache_size_bytes=100 * 1024 * 1024, # 100MB
)
purview_agent_middleware = PurviewPolicyMiddleware(build_credential(), settings)
agent = Agent(
client=client,
@@ -296,16 +284,8 @@ async def run_with_custom_cache_provider() -> None:
print("-- Default Cache Path --")
print("Using default InMemoryCacheProvider with settings-based configuration")
first: AgentResponse = await agent.run(
Message("user", ["Tell me a joke about AI."], additional_properties={"user_id": user_id})
)
print("First response (default cache):\n", first)
second: AgentResponse = await agent.run(
Message("user", ["Nice! Another AI joke please."], additional_properties={"user_id": user_id})
)
print("Second response (default cache):\n", second)
blocked_text = settings.get("blocked_prompt_message") or "Prompt blocked by policy"
await run_policy_flow("default cache", agent, user_id, blocked_text)
async def main() -> None:
@@ -326,6 +306,11 @@ async def main() -> None:
except Exception as ex: # pragma: no cover - demo resilience
print(f"Custom cache provider path failed: {ex}")
try:
await run_with_default_cache()
except Exception as ex: # pragma: no cover - demo resilience
print(f"Default cache path failed: {ex}")
if __name__ == "__main__":
asyncio.run(main())
python/uv.lock
Generated
+10 -10
View File
@@ -115,7 +115,7 @@ wheels = [
[[package]]
name = "agent-framework"
version = "1.8.0"
version = "1.8.1"
source = { virtual = "." }
dependencies = [
{ name = "agent-framework-core", extra = ["all"], marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -185,7 +185,7 @@ requires-dist = [
[[package]]
name = "agent-framework-ag-ui"
version = "1.0.0rc3"
version = "1.0.0rc4"
source = { editable = "packages/ag-ui" }
dependencies = [
{ name = "ag-ui-protocol", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -279,7 +279,7 @@ requires-dist = [
[[package]]
name = "agent-framework-azurefunctions"
version = "1.0.0b260604"
version = "1.0.0b260609"
source = { editable = "packages/azurefunctions" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -333,7 +333,7 @@ requires-dist = [
[[package]]
name = "agent-framework-claude"
version = "1.0.0b260521"
version = "1.0.0b260609"
source = { editable = "packages/claude" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -363,7 +363,7 @@ requires-dist = [
[[package]]
name = "agent-framework-core"
version = "1.8.0"
version = "1.8.1"
source = { editable = "packages/core" }
dependencies = [
{ name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -529,7 +529,7 @@ dev = [{ name = "types-python-dateutil", specifier = "==2.9.0.20260518" }]
[[package]]
name = "agent-framework-foundry"
version = "1.8.0"
version = "1.8.1"
source = { editable = "packages/foundry" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -548,7 +548,7 @@ requires-dist = [
[[package]]
name = "agent-framework-foundry-hosting"
version = "1.0.0a260604"
version = "1.0.0a260609"
source = { editable = "packages/foundry_hosting" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -584,7 +584,7 @@ requires-dist = [
[[package]]
name = "agent-framework-gemini"
version = "1.0.0a260521"
version = "1.0.0a260609"
source = { editable = "packages/gemini" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -714,7 +714,7 @@ dev = [
[[package]]
name = "agent-framework-mem0"
version = "1.0.0b260521"
version = "1.0.0b260609"
source = { editable = "packages/mem0" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -774,7 +774,7 @@ requires-dist = [
[[package]]
name = "agent-framework-openai"
version = "1.8.0"
version = "1.8.1"
source = { editable = "packages/openai" }
dependencies = [
{ name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },