From d90355fb8f5e63c4ce2753e3a546c11f65b9241c Mon Sep 17 00:00:00 2001 From: Lubaoshuai <128781758+Lubaoshuai@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:15:50 +0000 Subject: [PATCH 01/26] Python: exclude encrypted reasoning payloads from compaction token counts (#8354) * Python: exclude encrypted reasoning payloads from compaction token counts The compaction token counter serialises each message and tokenises the JSON. Content.protected_data (OpenAI encrypted_content, Anthropic thinking signature, Gemini thought_signature, OpenAI reasoning_details) is opaque ciphertext the model replays but never tokenises, so counting it inflated every token-aware compaction decision. The same blob can also ride in additional_properties as encrypted_content. Both are now excluded from serialised content, matching the existing exclusions for raw_representation and mirrored items. Fixes #8348 * fix(compaction): count clear-text reasoning details and invalidate stale token caches Review follow-up for the protected-data token exclusion: - reasoning_details stored in protected_data are JSON-serialised and replayed to the provider as clear text (summary, reasoning_text and nested reasoning.text), so only opaque members (encrypted_content) are excluded from the estimate; non-JSON payloads (Anthropic signature, Responses API encrypted blobs) are still dropped entirely. - group token counts are stamped with a serialization basis version; counts cached under an older basis (e.g. sessions annotated before a serialization change) are recomputed instead of reused. * fix(compaction): resolve pyright/mypy typing errors in token estimation - annotate the recursive reasoning-payload filter with explicit dict[Any, Any] / list[Any] casts so iterating the narrowed Any dict does not leak Unknown types under pyright strict - cast additional_properties before the encrypted_content filter - assert non-None token counts in the clear-text regression test so the ordering comparison does not involve optional integers --------- Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com> --- .../core/agent_framework/_compaction.py | 77 ++++++++++++- .../core/tests/core/test_compaction.py | 103 ++++++++++++++++++ 2 files changed, 179 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 227417bb57a..45205866561 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -30,6 +30,12 @@ GROUP_INDEX_KEY = "index" GROUP_HAS_REASONING_KEY = "has_reasoning" GROUP_TOKEN_COUNT_KEY = "token_count" # ruff:ignore[hardcoded-password-string] # nosec B105 - compaction metadata key, not a credential +GROUP_TOKEN_COUNT_BASIS_KEY = "token_count_basis" # ruff:ignore[hardcoded-password-string] # nosec B105 - compaction metadata key, not a credential +# Bumped whenever the token-estimation serialization changes. Cached counts +# stamped with an older basis are recomputed instead of reused: annotations +# survive Message.to_dict()/from_dict() round-trips, so without the basis a +# transcript annotated before a serialization change keeps its stale counts. +TOKEN_COUNT_BASIS_VERSION = 2 EXCLUDED_KEY = "_excluded" EXCLUDE_REASON_KEY = "_exclude_reason" SUMMARY_OF_MESSAGE_IDS_KEY = "_summary_of_message_ids" @@ -496,7 +502,13 @@ def _write_group_annotation( token_count: int | None = None if existing_raw_annotation is not None: raw_token_count = existing_raw_annotation.get(GROUP_TOKEN_COUNT_KEY) - if isinstance(raw_token_count, int) or raw_token_count is None: + raw_token_count_basis = existing_raw_annotation.get(GROUP_TOKEN_COUNT_BASIS_KEY) + # Preserve a cached count only when it was computed with the current + # serialization basis; counts from older bases are stale and get + # recomputed (see TOKEN_COUNT_BASIS_VERSION). + if raw_token_count_basis == TOKEN_COUNT_BASIS_VERSION and ( + isinstance(raw_token_count, int) or raw_token_count is None + ): token_count = raw_token_count unknown_fields = { key: value @@ -508,6 +520,7 @@ def _write_group_annotation( GROUP_INDEX_KEY, GROUP_HAS_REASONING_KEY, GROUP_TOKEN_COUNT_KEY, + GROUP_TOKEN_COUNT_BASIS_KEY, } } @@ -517,6 +530,7 @@ def _write_group_annotation( GROUP_INDEX_KEY: index, GROUP_HAS_REASONING_KEY: has_reasoning, GROUP_TOKEN_COUNT_KEY: token_count, + GROUP_TOKEN_COUNT_BASIS_KEY: TOKEN_COUNT_BASIS_VERSION if token_count is not None else None, } annotation.update(unknown_fields) message.additional_properties[GROUP_ANNOTATION_KEY] = annotation @@ -550,6 +564,10 @@ def _token_count(message: Message) -> int | None: if annotation is None: return None token_count = annotation.get(GROUP_TOKEN_COUNT_KEY) + # Counts stamped with an older serialization basis are stale: they were + # computed from a different serialized form and must be recomputed. + if annotation.get(GROUP_TOKEN_COUNT_BASIS_KEY) != TOKEN_COUNT_BASIS_VERSION: + return None return token_count if isinstance(token_count, int) else None @@ -558,6 +576,7 @@ def _write_token_count(message: Message, token_count: int) -> None: if annotation is None: return annotation[GROUP_TOKEN_COUNT_KEY] = token_count + annotation[GROUP_TOKEN_COUNT_BASIS_KEY] = TOKEN_COUNT_BASIS_VERSION message.additional_properties[GROUP_ANNOTATION_KEY] = annotation @@ -732,12 +751,68 @@ def annotate_message_groups( return _ordered_group_ids_from_annotations(messages) +_OPAQUE_REASONING_KEYS: Final[frozenset[str]] = frozenset({"encrypted_content"}) + + +def _strip_opaque_reasoning_payload(value: Any) -> Any: + """Recursively drop opaque reasoning members, keeping clear-text payloads.""" + if isinstance(value, dict): + entries = cast("dict[Any, Any]", value) + filtered: dict[Any, Any] = {} + for key, item in entries.items(): + if key in _OPAQUE_REASONING_KEYS: + continue + stripped = _strip_opaque_reasoning_payload(item) + if stripped is not None: + filtered[key] = stripped + return filtered or None + if isinstance(value, list): + entries = cast("list[Any]", value) + kept = [ + stripped + for stripped in (_strip_opaque_reasoning_payload(entry) for entry in entries) + if stripped is not None + ] + return kept or None + return value + + def _serialize_content(content: Content) -> dict[str, Any]: payload = content.to_dict(exclude_none=True) payload.pop("raw_representation", None) # ``items`` mirrors ``result`` for function_result content; exclude it # to avoid double-counting tokens during estimation. payload.pop("items", None) + # ``protected_data`` carries provider reasoning payloads. JSON-serialised + # reasoning_details (Chat Completions) are replayed to the provider as + # clear text -- ``summary``, ``reasoning_text`` and nested ``reasoning.text`` + # are part of the context the provider receives, so only their opaque + # members (``encrypted_content``) are excluded and the clear text stays + # counted. Anything that is not such a JSON structure (Anthropic thinking + # ``signature``, Responses API ``encrypted_content`` blobs) is replayed + # opaquely and never tokenised; exclude it so estimation measures the text + # the model actually sees. + protected_data = payload.get("protected_data") + if isinstance(protected_data, str) and protected_data: + try: + reasoning_payload = json.loads(protected_data) + except ValueError: + reasoning_payload = None + filtered = ( + _strip_opaque_reasoning_payload(reasoning_payload) if isinstance(reasoning_payload, (dict, list)) else None + ) + if filtered is None: + payload.pop("protected_data", None) + else: + payload["protected_data"] = json.dumps(filtered, ensure_ascii=False) + else: + payload.pop("protected_data", None) + additional_properties = payload.get("additional_properties") + if isinstance(additional_properties, dict) and "encrypted_content" in additional_properties: + typed_properties = cast("dict[str, Any]", additional_properties) + payload["additional_properties"] = { + key: value for key, value in typed_properties.items() if key != "encrypted_content" + } return payload diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 8e48faa4bf4..c02eb5364b4 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -435,7 +435,106 @@ def test_annotate_message_groups_with_tokenizer_adds_token_counts() -> None: assert isinstance(_token_count(messages[1]), int) +def test_annotate_message_groups_token_counts_ignore_protected_data() -> None: + encrypted_blob = "A" * 4096 + with_payload = Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_1", text="", protected_data=encrypted_blob), + Content.from_text("final answer"), + ], + ) + without_payload = Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_1", text=""), + Content.from_text("final answer"), + ], + ) + + annotate_message_groups([with_payload], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([without_payload], tokenizer=CharacterEstimatorTokenizer()) + + assert _token_count(with_payload) == _token_count(without_payload) + + +def test_annotate_message_groups_token_counts_ignore_encrypted_content_in_additional_properties() -> None: + encrypted_blob = "A" * 4096 + with_blob = Content.from_text_reasoning(id="rs_1", text="") + with_blob.additional_properties["encrypted_content"] = encrypted_blob + without_blob = Content.from_text_reasoning(id="rs_1", text="") + + with_message = Message(role="assistant", contents=[with_blob, Content.from_text("final answer")]) + without_message = Message(role="assistant", contents=[without_blob, Content.from_text("final answer")]) + + annotate_message_groups([with_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([without_message], tokenizer=CharacterEstimatorTokenizer()) + + assert _token_count(with_message) == _token_count(without_message) + + +def test_annotate_message_groups_token_counts_keep_clear_text_reasoning_details() -> None: + import json + + clear_text = "visible reasoning " * 128 + with_details = Content.from_text_reasoning( + id="rs_1", + text="", + protected_data=json.dumps([{"type": "reasoning.text", "text": clear_text, "encrypted_content": "A" * 4096}]), + ) + summary_only = Content.from_text_reasoning( + id="rs_1", + text="", + protected_data=json.dumps([{"type": "reasoning.text", "text": clear_text}]), + ) + without_details = Content.from_text_reasoning(id="rs_1", text="") + final_answer = Content.from_text("final answer") + + with_message = Message(role="assistant", contents=[with_details, final_answer]) + summary_only_message = Message(role="assistant", contents=[summary_only, final_answer]) + bare_message = Message(role="assistant", contents=[without_details, final_answer]) + + annotate_message_groups([with_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([summary_only_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([bare_message], tokenizer=CharacterEstimatorTokenizer()) + + # The opaque encrypted member is excluded while the clear-text reasoning + # the provider replays stays counted. + with_count = _token_count(with_message) + summary_count = _token_count(summary_only_message) + bare_count = _token_count(bare_message) + assert with_count is not None + assert summary_count is not None + assert bare_count is not None + assert with_count == summary_count + assert with_count > bare_count + + +def test_annotate_token_counts_recomputes_counts_from_stale_serialization_basis() -> None: + + from agent_framework._compaction import TOKEN_COUNT_BASIS_VERSION, _write_group_annotation, annotate_token_counts + + message = Message(role="assistant", contents=[Content.from_text("final answer " * 64)]) + _write_group_annotation(message, group_id="g0", kind="assistant_text", index=0, has_reasoning=False) + annotation = message.additional_properties[GROUP_ANNOTATION_KEY] + assert isinstance(annotation, dict) + # Simulate a count cached by a pre-fix serialization basis: value present, + # no basis stamp (annotations survive to_dict/from_dict round-trips). + annotation[GROUP_TOKEN_COUNT_KEY] = 1 + annotation.pop("token_count_basis", None) + assert _token_count(message) == 1 + + annotate_token_counts([message], tokenizer=CharacterEstimatorTokenizer()) + + annotation = message.additional_properties[GROUP_ANNOTATION_KEY] + assert isinstance(annotation, dict) + assert annotation[GROUP_TOKEN_COUNT_KEY] != 1 + assert annotation["token_count_basis"] == TOKEN_COUNT_BASIS_VERSION + + def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() -> None: + from agent_framework._compaction import TOKEN_COUNT_BASIS_VERSION + tokenizer = CharacterEstimatorTokenizer() messages = [_assistant_function_call("c3")] annotate_message_groups(messages) @@ -445,6 +544,10 @@ def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() annotation = messages[0].additional_properties.get(GROUP_ANNOTATION_KEY) if isinstance(annotation, dict): annotation[GROUP_TOKEN_COUNT_KEY] = old_token_count + # Counts cached under the current serialization basis are preserved + # across extensions; counts without the current basis stamp are stale + # and get recomputed (see test_annotate_token_counts_recomputes_...). + annotation["token_count_basis"] = TOKEN_COUNT_BASIS_VERSION extend_compaction_messages(messages, [_tool_result("c3", "ok")], tokenizer=tokenizer) From f187b688263084f09c6c682e145d332de448be35 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Thu, 17 Sep 2026 08:19:06 +0000 Subject: [PATCH 02/26] Python: preserve approval context across resume paths (#8433) * Python: fix mixed function-call batch classification * Python: scope mixed pause correlation * Python: preserve streaming approval updates * Python: preserve approval resume context * Python: address approval resume review feedback * Python: harden mixed approval replay * Python: keep stateless mixed approvals inert --- .../specs/004-python-function-calling-loop.md | 13 +- python/packages/ag-ui/AGENTS.md | 4 + .../ag-ui/agent_framework_ag_ui/_agent.py | 6 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 24 +- .../ag_ui/test_agent_wrapper_comprehensive.py | 87 +++- .../tests/ag_ui/test_approval_result_event.py | 39 +- .../packages/core/agent_framework/_agents.py | 6 + .../_harness/_tool_approval.py | 17 +- .../packages/core/agent_framework/_tools.py | 476 +++++++++--------- .../core/test_function_invocation_logic.py | 237 +++++++++ .../tests/core/test_harness_tool_approval.py | 180 +++++++ 11 files changed, 844 insertions(+), 245 deletions(-) diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index bfbe4812299..4603417996c 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -191,9 +191,12 @@ Before acting on a model function-call batch, the loop classifies every actionab call aborts the complete batch before approval state changes or execution. Otherwise approval-required and Host-owned calls are returned together in model order, while session-backed executable siblings remain deferred. An incomplete session-backed mixed approval/Host response remains pending without executing a deferred call; a stateless incomplete -response is rejected. Correlation is scoped to the active mixed batch so completed or abandoned historical Host calls -remain unchanged. `ToolApprovalMiddleware` may resolve approval requests through standing or automatic policies, but -it preserves non-approval user-input requests and does not split manual approvals away from their Host-owned siblings. +response is rejected, including when no response is supplied, message roles vary within the model output, or a newer +standalone request follows the incomplete batch. Stateless discovery selects the latest unresolved mixed batch, so a +completed batch remains inert on later turns. Correlation is scoped to the active mixed batch so completed or abandoned +historical Host calls remain unchanged. +`ToolApprovalMiddleware` may resolve approval requests through standing or automatic policies, but it preserves +non-approval user-input requests and does not split or reorder manual approvals relative to their Host-owned siblings. ### Reasoning-bound function-call groups @@ -612,7 +615,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Scenario | Required invariant | Primary regression test | |---|---|---| | Fatal call mixed with pauses | Complete-batch classification raises before approval or execution, independent of call order. | `packages/core/tests/core/test_function_invocation_logic.py::test_mixed_batch_fatal_unknown_precedes_every_pause` | -| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_active_mixed_pause_ignores_historical_host_requests` | +| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; a stateless zero-response or partial response fails closed across message roles and cannot be hidden by a newer standalone request; completed mixed batches remain inert; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution`, `test_stateless_mixed_batch_across_message_roles_requires_complete_responses`, `test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch`, `test_completed_split_stateless_mixed_batch_is_inert_on_later_turn`, `test_active_mixed_pause_ignores_historical_host_requests` | | Safe and approval-required calls in one batch | Hidden safe calls replay only with the matching visible approval. | `packages/core/tests/core/test_harness_tool_approval.py::test_mixed_batch_hides_already_approved_request_until_approval_replay` | | Restored approval state | Serialized `ToolApprovalState` restores mixed-batch behavior. | `test_mixed_batch_accepts_restored_tool_approval_state` | | Unrelated turn before approval | Hidden calls do not execute on an unrelated turn. | `test_hidden_mixed_batch_requests_do_not_replay_on_unrelated_turn` | @@ -620,7 +623,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Queued approvals | One unresolved approval is surfaced per run without premature execution. | `test_tool_approval_middleware_queues_multiple_approval_requests`, `test_tool_approval_middleware_queues_streamed_approval_requests` | | Middleware state plus hidden core state | State saves do not discard hidden mixed-batch calls. | `test_tool_approval_middleware_preserves_hidden_mixed_batch_requests` | | Auto-approval callback | Callback receives the original function call and executes the approved set once. | `test_tool_approval_middleware_auto_approval_rule_receives_function_call` | -| Approval policy with Host-owned sibling | Reordering the same calls has the same outcome; auto-approved and safe calls remain deferred until Host input arrives; manual approvals stay in the mixed batch; an approved Host tool returns to Host-owned handling before local execution. | `test_tool_approval_middleware_mixed_batch_is_order_independent`, `test_tool_approval_middleware_auto_approves_with_host_pause_and_cached_safe_call`, `test_tool_approval_middleware_keeps_manual_approvals_together_with_host_pause`, `test_tool_approval_middleware_policy_approval_reclassifies_host_tool` | +| Approval policy with Host-owned sibling | Reordering the same calls has the same outcome; auto-approved and safe calls remain deferred until Host input arrives; separately streamed Host input stays immediately visible while later manual approvals remain in model order and in the same batch; an approved Host tool returns to Host-owned handling before local execution. | `test_tool_approval_middleware_mixed_batch_is_order_independent`, `test_tool_approval_middleware_auto_approves_with_host_pause_and_cached_safe_call`, `test_tool_approval_middleware_keeps_manual_approvals_together_with_host_pause`, `test_tool_approval_middleware_preserves_split_streamed_mixed_batch_order`, `test_tool_approval_middleware_streams_standalone_user_input_immediately`, `test_tool_approval_middleware_policy_approval_reclassifies_host_tool` | | Shared call budget | Auto-approved re-entry does not reset `max_function_calls`, and every executed approval group counts even when it pauses for input. | `test_tool_approval_middleware_auto_approved_loops_share_function_call_budget`, `test_approval_resume_user_input_counts_toward_function_call_budget` | | Standing tool rule | Tool-level approval applies only to later matching tools. | `test_tool_approval_middleware_always_approve_tool_rule` | | Forged standing rule | An unbound or substituted hosted response cannot create a standing middleware approval rule for caller-selected metadata. | `test_tool_approval_middleware_drops_forged_standing_approval`, `test_tool_approval_middleware_rebinds_hosted_standing_approval` | diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 25233eaf56d..94a11e1ca29 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -41,6 +41,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. `Agent.run` path, including Agent middleware, provider `before_run` preparation, and function middleware. `_approval_execution.py` tracks execution through existing function/chat middleware interfaces; it does not prepare providers or invoke the tools itself. `MiddlewareFailure` remains fatal. +- Programmatic `AgentFrameworkAgent.run(...)` and `run_agent_stream(...)` calls accept + `function_invocation_kwargs` and preserve that tool-runtime context across approval resume. Adapter-local execution + merges only that mapping, the wrapped agent's default `additional_function_arguments`, and the current run options + (in increasing precedence), not unrelated AG-UI run kwargs. - Queued approvals follow the wrapped Agent's scheduling. With `ToolApprovalMiddleware`, collected decisions remain pending until that middleware releases the batch; AG-UI does not execute a tool ahead of it. Collected server-side grants do not require the client to submit the same approval again. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 9d127f2bbff..ab179da165c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -3,7 +3,7 @@ """AgentFrameworkAgent wrapper for AG-UI protocol.""" import warnings -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from typing import Any, cast from ag_ui.core import BaseEvent @@ -175,11 +175,14 @@ def snapshot_store(self) -> AGUIThreadSnapshotStore | None: async def run( self, input_data: dict[str, Any], + *, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AsyncGenerator[BaseEvent, None]: """Run the wrapped agent and yield AG-UI events. Args: input_data: The AG-UI run input containing messages, state, etc. + function_invocation_kwargs: Keyword arguments forwarded only to tool invocation. Yields: AG-UI events @@ -190,5 +193,6 @@ async def run( self.agent, self.config, approval_state_store=self._approval_state_store, + function_invocation_kwargs=function_invocation_kwargs, ): yield event diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 825a60ec621..1baf001c640 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -1938,7 +1938,22 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] raise TypeError("Local approval execution with context providers requires a core Agent execution path.") client = getattr(agent, "client", None) config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None)) - tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"} + raw_function_invocation_kwargs = run_kwargs.get("function_invocation_kwargs") + tool_kwargs = ( + dict(cast(Mapping[str, Any], raw_function_invocation_kwargs)) + if raw_function_invocation_kwargs is not None + else {} + ) + default_options = getattr(agent, "default_options", None) + if isinstance(default_options, Mapping) and ( + default_function_arguments := default_options.get("additional_function_arguments") + ): + tool_kwargs.update(cast(Mapping[str, Any], default_function_arguments)) + raw_options = run_kwargs.get("options") + if isinstance(raw_options, Mapping) and ( + additional_function_arguments := raw_options.get("additional_function_arguments") + ): + tool_kwargs.update(cast(Mapping[str, Any], additional_function_arguments)) for approval in static_approved: function_call = approval.function_call call_id = (function_call.call_id if function_call else None) or approval.id or "" @@ -2668,6 +2683,8 @@ async def run_agent_stream( agent: SupportsAgentRun, config: AgentConfig, approval_state_store: InMemoryAGUIApprovalStateStore | None = None, + *, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AsyncGenerator[BaseEvent]: """Run agent and yield AG-UI events. @@ -2680,6 +2697,7 @@ async def run_agent_stream( config: Agent configuration approval_state_store: Optional server-side Approval State store used to preserve approval-only middleware state across AG-UI requests. + function_invocation_kwargs: Keyword arguments forwarded only to tool invocation. Yields: AG-UI events @@ -2692,6 +2710,7 @@ async def run_agent_stream( config, state_store, authorized_executions=authorized_executions, + function_invocation_kwargs=function_invocation_kwargs, ) try: async for event in stream: @@ -2711,6 +2730,7 @@ async def _run_agent_stream( approval_state_store: InMemoryAGUIApprovalStateStore, *, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution], + function_invocation_kwargs: Mapping[str, Any] | None, ) -> AsyncGenerator[BaseEvent]: # Parse IDs supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") @@ -3105,6 +3125,8 @@ async def _run_agent_stream( # Build run kwargs (Feature #6: Azure store flag when metadata present) run_kwargs: dict[str, Any] = {"session": session} + if function_invocation_kwargs is not None: + run_kwargs["function_invocation_kwargs"] = dict(function_invocation_kwargs) if tools: run_kwargs["tools"] = tools # Hand the forwarded AG-UI context to the A2UI runner PER REQUEST (not just at diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index b32cea9af8c..2a01c03a4b4 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -8,7 +8,16 @@ import pytest from ag_ui.core import RunErrorEvent, ToolCallResultEvent -from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message +from agent_framework import ( + Agent, + ChatOptions, + ChatResponseUpdate, + Content, + ContextProvider, + FunctionInvocationContext, + Message, + tool, +) from pydantic import BaseModel from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner @@ -23,6 +32,82 @@ def _approval_request_id(events: list[Any]) -> str: return str(event.value["id"]) +async def test_approval_resume_with_context_provider_preserves_function_invocation_kwargs( + streaming_chat_client_stub: Any, +) -> None: + """Context-provider preparation and runtime tool kwargs survive an AG-UI resume.""" + from agent_framework.ag_ui import AgentFrameworkAgent + + observed_kwargs: list[dict[str, Any]] = [] + phase = "pause" + + class NoOpContextProvider(ContextProvider): + pass + + @tool(name="sensitive_action", approval_mode="always_require") + def sensitive_action(context: FunctionInvocationContext) -> str: + observed_kwargs.append(dict(context.kwargs)) + return "executed" + + async def stream_fn( + messages: MutableSequence[Message], + options: ChatOptions, + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + if phase == "pause": + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + name="sensitive_action", + call_id="call-sensitive", + arguments="{}", + ) + ] + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Done")]) + + wrapper = AgentFrameworkAgent( + agent=Agent( + client=streaming_chat_client_stub(stream_fn), + name="test_agent", + instructions="Test", + tools=[sensitive_action], + context_providers=[NoOpContextProvider("context")], + ) + ) + runtime_kwargs = {"user_id": "user-123"} + thread_id = "context-provider-runtime-kwargs" + pause_events = [ + event + async for event in wrapper.run( + {"thread_id": thread_id, "messages": [{"role": "user", "content": "do it"}]}, + function_invocation_kwargs=runtime_kwargs, + ) + ] + approval_id = _approval_request_id(pause_events) + + phase = "resume" + resume_events = [ + event + async for event in wrapper.run( + { + "thread_id": thread_id, + "messages": [], + "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], + }, + function_invocation_kwargs=runtime_kwargs, + ) + ] + + assert observed_kwargs[0]["user_id"] == runtime_kwargs["user_id"] + assert set(observed_kwargs[0]) == {"session", "user_id"} + assert [ + (event.tool_call_id, event.content) for event in resume_events if isinstance(event, ToolCallResultEvent) + ] == [("call-sensitive", "executed")] + + async def test_agent_initialization_basic(streaming_chat_client_stub): """Test basic agent initialization without state schema.""" from agent_framework.ag_ui import AgentFrameworkAgent diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index a43dde00d19..86b52f068aa 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -8,7 +8,7 @@ from typing import Any from ag_ui.core import ToolCallResultEvent -from agent_framework import AgentResponseUpdate, Content, FunctionTool +from agent_framework import AgentResponseUpdate, Content, FunctionInvocationContext, FunctionTool from agent_framework.exceptions import UserInputRequiredException from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] @@ -88,10 +88,18 @@ async def _run_resume( return events -async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent]: +async def _run_custom_approval( + tool: FunctionTool, + *, + function_invocation_kwargs: dict[str, Any] | None = None, + default_additional_function_arguments: dict[str, Any] | None = None, +) -> tuple[list[Any], StubAgent]: agent = StubAgent( updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, + default_options={ + "tools": [tool], + "additional_function_arguments": default_additional_function_arguments, + }, ) store = InMemoryAGUIApprovalStateStore() store.register( @@ -114,6 +122,7 @@ async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent agent, AgentConfig(), approval_state_store=store, + function_invocation_kwargs=function_invocation_kwargs, ) ] return events, agent @@ -134,6 +143,30 @@ async def test_approved_call_emits_one_live_result_under_original_identity() -> assert [(event.tool_call_id, event.content) for event in results] == [("call-weather", "Sunny in Seattle")] +async def test_approved_call_merges_function_invocation_kwargs_with_agent_defaults() -> None: + """Static approval execution uses normal runtime/default precedence without run-kwarg leakage.""" + observed_kwargs: list[dict[str, Any]] = [] + + def inspect_context(context: FunctionInvocationContext) -> str: + observed_kwargs.append(dict(context.kwargs)) + return "inspected" + + events, _ = await _run_custom_approval( + FunctionTool(name="inspect_context", description="Inspect context", func=inspect_context), + function_invocation_kwargs={"runtime_only": "runtime", "shared": "runtime"}, + default_additional_function_arguments={"default_only": "default", "shared": "default"}, + ) + + results = [event for event in events if isinstance(event, ToolCallResultEvent)] + assert observed_kwargs[0] == { + "default_only": "default", + "runtime_only": "runtime", + "session": observed_kwargs[0]["session"], + "shared": "default", + } + assert [(event.tool_call_id, event.content) for event in results] == [("call-custom", "inspected")] + + async def test_rejected_call_does_not_execute_or_emit_live_result() -> None: executions: list[str] = [] diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c4f5da1c0de..205d44b35c6 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1405,8 +1405,10 @@ async def _prepare_run_context( self.context_providers.append(InMemoryHistoryProvider()) active_session = session + framework_created_session = False if active_session is None and self.context_providers: active_session = AgentSession() + framework_created_session = True per_service_call_history_providers = self._resolve_per_service_call_history_providers( session=active_session, @@ -1529,8 +1531,12 @@ async def _prepare_run_context( session_messages: list[Message] = session_context.get_messages(include_input=True) effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {} + from ._tools import _APPROVAL_SESSION_IS_AUTHORITATIVE_KEY # pyright: ignore[reportPrivateUsage] + + effective_client_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, None) if active_session is not None: effective_client_kwargs["session"] = active_session + effective_client_kwargs[_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY] = not framework_created_session per_service_call_history_middleware: PerServiceCallHistoryPersistingMiddleware | None = None if per_service_call_history_providers and active_session is not None: per_service_call_history_middleware = PerServiceCallHistoryPersistingMiddleware( diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index ad81d06451d..4e2db951c77 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -458,14 +458,22 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: approval_requests: list[Content] = [] buffered_approval_updates: list[AgentResponseUpdate] = [] + streamed_user_input_contents: list[Content] = [] + buffering = False async for update in context.result: update_approval_requests = [ content for content in update.contents if content.type == "function_approval_request" ] - if not update_approval_requests: + if not buffering and not update_approval_requests: + streamed_user_input_contents.extend( + content for content in update.contents if content.user_input_request + ) + if update_approval_requests: + buffering = True + approval_requests.extend(update_approval_requests) + if not buffering: yield update continue - approval_requests.extend(update_approval_requests) buffered_update = copy.copy(update) buffered_update.contents = list(update.contents) buffered_approval_updates.append(buffered_update) @@ -476,7 +484,10 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: response_messages = [ Message( role="assistant", - contents=[content for update in buffered_approval_updates for content in update.contents], + contents=[ + *streamed_user_input_contents, + *(content for update in buffered_approval_updates for content in update.contents), + ], ) ] has_other_user_input = self._has_non_approval_user_input(response_messages) diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 1caa4d82565..a292383ac7a 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -105,15 +105,12 @@ def _generate_function_call_occurrence_id() -> str: DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3 SHELL_TOOL_KIND_VALUE: Final[str] = "shell" _TOOL_APPROVAL_STATE_KEY: Final[str] = "tool_approval" -_RUN_LOCAL_MIDDLEWARE_SESSION_ATTR: Final[str] = "_run_local_function_middleware_session" +_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY: Final[str] = "_approval_session_is_authoritative" def _has_authoritative_approval_session(invocation_session: AgentSession | None) -> bool: """Return whether approval state belongs to a caller-owned session.""" - return ( - invocation_session is not None - and getattr(invocation_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, False) is not True - ) + return invocation_session is not None _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY: Final[str] = "already_approved_approval_request_groups" @@ -2257,6 +2254,7 @@ async def _try_execute_function_call_groups( tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]], config: FunctionInvocationConfiguration, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, middleware_pipeline: FunctionMiddlewarePipeline | None = None, host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> tuple[list[list[Content]], bool]: @@ -2268,6 +2266,7 @@ async def _try_execute_function_call_groups( tools: The tools available for execution. config: Configuration for function invocation. invocation_session: The agent session for this invocation, if any. + approval_session_is_authoritative: Whether the invocation session owns persistent approval state. middleware_pipeline: Optional middleware pipeline to apply during execution. host_payload_budget: Shared request budget for retained Host-only function result payloads. @@ -2278,6 +2277,8 @@ async def _try_execute_function_call_groups( """ from ._types import Content + approval_session = invocation_session if approval_session_is_authoritative else None + # Normalize the batch to calls owned by this layer before making any control-flow decision. function_calls = [ function_call @@ -2363,18 +2364,18 @@ async def _try_execute_function_call_groups( visible_requests.append(approval_request) pause_groups.append([approval_request]) continue - if not _has_authoritative_approval_session(invocation_session): + if not _has_authoritative_approval_session(approval_session): visible_requests.append(approval_request) pause_groups.append([approval_request]) continue already_approved_requests.append(approval_request) _store_already_approved_approval_requests( - invocation_session, + approval_session, visible_requests, already_approved_requests, ) - _store_pending_approval_requests(invocation_session, visible_requests) - _store_pending_mixed_pause_batch(invocation_session, pause_groups) + _store_pending_approval_requests(approval_session, visible_requests) + _store_pending_mixed_pause_batch(approval_session, pause_groups) return pause_groups, False if has_declaration_only_call: # Declaration-only calls are returned as user input rather than executed locally. @@ -2470,6 +2471,7 @@ async def _execute_function_calls( options: dict[str, Any] | None, config: FunctionInvocationConfiguration, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, middleware_pipeline: FunctionMiddlewarePipeline | None = None, host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> _FunctionExecutionBatch: @@ -2481,6 +2483,7 @@ async def _execute_function_calls( function_calls=function_calls, tools=tools, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=middleware_pipeline, config=config, host_payload_budget=host_payload_budget, @@ -2966,6 +2969,110 @@ def _same_mixed_pause_response(left: Mapping[str, Any], right: Mapping[str, Any] return left_payload == right_payload +def _match_mixed_pause_responses( + items: list[dict[str, Any]], + responses: Sequence[Content], + *, + approval_response_binder: Callable[[Content], Content | None] | None = None, +) -> tuple[set[int], bool, list[Content], set[int]]: + """Match one complete mixed pause batch without depending on its storage source.""" + approval_items: dict[str, int | None] = {} + host_items_by_occurrence: dict[str, int] = {} + host_items_by_call: dict[str, list[int]] = {} + for index, item in enumerate(items): + request = _content_from_state(item.get("request")) + if request is None: + continue + if item.get("kind") == "approval": + for identity in ( + request.id, + request.function_call.id if request.function_call is not None else None, + ): + if identity is None: + continue + if identity not in approval_items: + approval_items[identity] = index + elif approval_items[identity] != index: + approval_items[identity] = None + elif item.get("kind") == "host" and request.call_id is not None: + host_items_by_call.setdefault(request.call_id, []).append(index) + if request.id is not None: + host_items_by_occurrence[request.id] = index + + matched_content_ids: set[int] = set() + for response in responses: + item_index: int | None = None + if response.type == "function_approval_response": + candidate = approval_response_binder(response) if approval_response_binder is not None else response + if candidate is None: + continue + response_identities = { + str(identity) + for identity in ( + candidate.additional_properties.get(_APPROVAL_REQUEST_ID_KEY), + candidate.id, + ) + if identity is not None + } + matching_indexes = { + matched_index + for identity in response_identities + if (matched_index := approval_items.get(identity)) is not None + } + if len(matching_indexes) == 1: + item_index = matching_indexes.pop() + elif response.type == "function_result" and response.call_id in host_items_by_call: + candidate = response + if response.id is not None: + item_index = host_items_by_occurrence.get(response.id) + request = _content_from_state(items[item_index].get("request")) if item_index is not None else None + if request is None or request.call_id != response.call_id: + item_index = None + else: + matching_indexes = [ + pending_index + for pending_index in host_items_by_call[response.call_id] + if items[pending_index].get("response") is None + or ( + isinstance(items[pending_index].get("response"), Mapping) + and _same_mixed_pause_response( + cast(Mapping[str, Any], items[pending_index]["response"]), + response.to_dict(), + ) + ) + ] + if len(matching_indexes) == 1: + item_index = matching_indexes[0] + else: + continue + + if item_index is None: + continue + candidate_state = candidate.to_dict() + stored_response = items[item_index].get("response") + if stored_response is not None and ( + not isinstance(stored_response, Mapping) + or not _same_mixed_pause_response(cast(Mapping[str, Any], stored_response), candidate_state) + ): + raise RuntimeError(f"Conflicting response for mixed pause occurrence {candidate.id!r}.") + items[item_index]["response"] = candidate_state + matched_content_ids.add(id(response)) + + if any(item.get("response") is None for item in items): + return matched_content_ids, True, [], set() + + ordered_responses: list[Content] = [] + host_result_ids: set[int] = set() + for item in items: + response = _content_from_state(item.get("response")) + if response is None: + return matched_content_ids, True, [], set() + ordered_responses.append(response) + if item.get("kind") == "host": + host_result_ids.add(id(response)) + return matched_content_ids, False, ordered_responses, host_result_ids + + def _stage_pending_mixed_pause_responses( messages: list[Message], invocation_session: AgentSession | None, @@ -2989,78 +3096,24 @@ def _stage_pending_mixed_pause_responses( state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) return False, False, set() - approval_items: dict[str, dict[str, Any]] = {} - host_items_by_occurrence: dict[str, dict[str, Any]] = {} - host_items_by_call: dict[str, list[dict[str, Any]]] = {} - for item in items: - request = _content_from_state(item.get("request")) - if request is None: - continue - if item.get("kind") == "approval": - for identity in ( - request.id, - request.function_call.id if request.function_call is not None else None, - ): - if identity is not None: - approval_items[identity] = item - elif item.get("kind") == "host" and request.call_id is not None: - host_items_by_call.setdefault(request.call_id, []).append(item) - if request.id is not None: - host_items_by_occurrence[request.id] = item - - matched_content_ids: set[int] = set() - for message in messages: - for content in message.contents: - item: dict[str, Any] | None = None - if content.type == "function_approval_response": - rebound = _bind_approval_response_to_pending_request( - content, - invocation_session, - consume=False, - ) - if rebound is None: - continue - request_id = rebound.additional_properties.get(_APPROVAL_REQUEST_ID_KEY) - item = approval_items.get(str(request_id)) if request_id is not None else None - if item is None and rebound.id is not None: - item = approval_items.get(rebound.id) - candidate = rebound - elif content.type == "function_result" and content.call_id in host_items_by_call: - candidate = content - if content.id is not None: - item = host_items_by_occurrence.get(content.id) - request = _content_from_state(item.get("request")) if item is not None else None - if request is None or request.call_id != content.call_id: - item = None - else: - matching_items = [ - pending_item - for pending_item in host_items_by_call[content.call_id] - if pending_item.get("response") is None - or ( - isinstance(pending_item.get("response"), Mapping) - and _same_mixed_pause_response( - cast(Mapping[str, Any], pending_item["response"]), - content.to_dict(), - ) - ) - ] - if len(matching_items) == 1: - item = matching_items[0] - else: - continue + def bind_approval_response(response: Content) -> Content | None: + return _bind_approval_response_to_pending_request( + response, + invocation_session, + consume=False, + ) - if item is None: - continue - candidate_state = candidate.to_dict() - stored_response = item.get("response") - if stored_response is not None and ( - not isinstance(stored_response, Mapping) - or not _same_mixed_pause_response(cast(Mapping[str, Any], stored_response), candidate_state) - ): - raise RuntimeError(f"Conflicting response for mixed pause occurrence {candidate.id!r}.") - item["response"] = candidate_state - matched_content_ids.add(id(content)) + responses = [ + content + for message in messages + for content in message.contents + if content.type in {"function_approval_response", "function_result"} + ] + matched_content_ids, incomplete, ordered_responses, host_result_ids = _match_mixed_pause_responses( + items, + responses, + approval_response_binder=bind_approval_response, + ) if matched_content_ids: filtered_messages: list[Message] = [] @@ -3071,18 +3124,9 @@ def _stage_pending_mixed_pause_responses( messages[:] = filtered_messages state[_PENDING_MIXED_PAUSE_BATCH_KEY] = {"items": items} - if any(item.get("response") is None for item in items): + if incomplete: return True, False, set() - ordered_responses: list[Content] = [] - host_result_ids: set[int] = set() - for item in items: - response = _content_from_state(item.get("response")) - if response is None: - return True, False, set() - ordered_responses.append(response) - if item.get("kind") == "host": - host_result_ids.add(id(response)) messages.append(Message(role="user", contents=ordered_responses)) return False, True, host_result_ids @@ -3090,145 +3134,99 @@ def _stage_pending_mixed_pause_responses( def _stateless_mixed_pause_batch_status( messages: list[Message], ) -> tuple[bool, set[int]]: - """Validate only the latest stateless mixed batch and order its responses.""" + """Validate the latest unresolved stateless mixed batch and order its responses.""" from ._types import Message - approval_requests: list[Content] = [] - host_requests: list[Content] = [] - approval_responses: list[Content] = [] - host_responses: list[Content] = [] - pause_order_contents: list[Content] = [] - batch_index: int | None = None - for index in range(len(messages) - 1, -1, -1): - message = messages[index] - if any(content.type == "function_approval_request" for content in message.contents) and any( - content.type == "function_call" and content.user_input_request for content in message.contents - ): - batch_index = index - break - if batch_index is None: - return False, set() - - for content in messages[batch_index].contents: - if content.type == "function_call": - pause_order_contents.append(content) - if content.user_input_request: - host_requests.append(content) - elif content.type == "function_approval_request": - pause_order_contents.append(content) - approval_requests.append(content) - for message in messages[batch_index + 1 :]: - for content in message.contents: - if content.type == "function_approval_response": - approval_responses.append(content) - elif content.type == "function_result": - host_responses.append(content) + flattened_contents = [content for message in messages for content in message.contents] + request_batches: list[tuple[int, list[dict[str, Any]], set[str]]] = [] + batch_items: list[dict[str, Any]] = [] + batch_kinds: set[str] = set() + last_request_index = -1 - if not approval_requests or not host_requests or not (approval_responses or host_responses): - return False, set() + def answers_current_batch(content: Content) -> bool: + if content.type not in {"function_approval_response", "function_result"}: + return False + for item in batch_items: + request = _content_from_state(item.get("request")) + if request is None: + continue + if item.get("kind") == "approval": + function_call = request.function_call + if content.type == "function_approval_response" and content.id in { + request.id, + function_call.id if function_call is not None else None, + }: + return True + if ( + content.type == "function_result" + and function_call is not None + and content.call_id == function_call.call_id + ): + return True + elif ( + item.get("kind") == "host" + and content.type == "function_result" + and content.call_id == request.call_id + and (content.id is None or request.id is None or content.id == request.id) + ): + return True + return False - approval_request_identities: list[set[str]] = [] - identity_owners: dict[str, int] = {} - for index, request in enumerate(approval_requests): - identities = { - identity - for identity in ( - request.id, - request.function_call.id if request.function_call is not None else None, - ) - if identity is not None - } - for identity in identities: - if identity in identity_owners and identity_owners[identity] != index: - return True, set() - identity_owners[identity] = index - approval_request_identities.append(identities) - - unmatched_approval_requests = set(range(len(approval_requests))) - matched_approval_responses: dict[int, Content] = {} - matched_approval_response_ids: set[int] = set() - for response in approval_responses: - if response.id is None: + for content_index, content in enumerate(flattened_contents): + if batch_items and answers_current_batch(content): + request_batches.append((last_request_index, batch_items, batch_kinds)) + batch_items = [] + batch_kinds = set() + if content.type == "function_approval_request": + kind = "approval" + elif content.type == "function_call" and content.user_input_request: + kind = "host" + else: continue - matching_indexes = [ - index for index, identities in enumerate(approval_request_identities) if response.id in identities - ] - if len(matching_indexes) != 1: + last_request_index = content_index + batch_kinds.add(kind) + batch_items.append({"kind": kind, "request": content.to_dict()}) + if batch_items: + request_batches.append((last_request_index, batch_items, batch_kinds)) + + for batch_end, items, kinds in reversed(request_batches): + if kinds != {"approval", "host"}: continue - matching_index = matching_indexes[0] - previous_response = matched_approval_responses.get(matching_index) - if previous_response is not None and previous_response.to_dict() != response.to_dict(): - function_call = approval_requests[matching_index].function_call - occurrence_id = function_call.id if function_call is not None else approval_requests[matching_index].id - raise RuntimeError(f"Conflicting approval response for occurrence {occurrence_id!r}.") - matched_approval_responses[matching_index] = response - unmatched_approval_requests.discard(matching_index) - matched_approval_response_ids.add(id(response)) - - unmatched_host_requests = list(host_requests) - matched_host_responses: dict[int, Content] = {} - matched_host_result_ids: set[int] = set() - for response in (candidate for candidate in host_responses if candidate.id is not None): - matching_index = next( - ( - index - for index, request in enumerate(unmatched_host_requests) - if request.id == response.id and request.call_id == response.call_id - ), - None, - ) - if matching_index is not None: - request = unmatched_host_requests.pop(matching_index) - matched_host_responses[id(request)] = response - matched_host_result_ids.add(id(response)) - - for response in (candidate for candidate in host_responses if candidate.id is None): - matching_indexes = [ - index for index, request in enumerate(unmatched_host_requests) if request.call_id == response.call_id + responses = [ + content + for content in flattened_contents[batch_end + 1 :] + if content.type in {"function_approval_response", "function_result"} ] - if len(matching_indexes) == 1: - request = unmatched_host_requests.pop(matching_indexes[0]) - matched_host_responses[id(request)] = response - matched_host_result_ids.add(id(response)) - - if unmatched_approval_requests or unmatched_host_requests: - return True, matched_host_result_ids - - approval_indexes_by_occurrence_id = { - request.function_call.id: index - for index, request in enumerate(approval_requests) - if request.function_call is not None and request.function_call.id is not None - } - approval_indexes_by_request_object = {id(request): index for index, request in enumerate(approval_requests)} - ordered_responses: list[Content] = [] - ordered_approval_indexes: set[int] = set() - for content in pause_order_contents: - host_response = matched_host_responses.get(id(content)) - if host_response is not None: - ordered_responses.append(host_response) + matched_response_ids, incomplete, ordered_responses, host_result_ids = _match_mixed_pause_responses( + items, + responses, + ) + if incomplete: + return True, host_result_ids + + pending_approval_response_ids = { + id(response) + for response in _collect_approval_responses( + messages, + non_approval_result_ids={ + id(content) + for content in responses + if content.type == "function_result" and id(content) in matched_response_ids + }, + ).values() + } + if matched_response_ids.isdisjoint(pending_approval_response_ids): continue - approval_index = approval_indexes_by_request_object.get(id(content)) - if approval_index is None and content.type == "function_call" and content.id is not None: - approval_index = approval_indexes_by_occurrence_id.get(content.id) - if approval_index is not None and approval_index not in ordered_approval_indexes: - ordered_responses.append(matched_approval_responses[approval_index]) - ordered_approval_indexes.add(approval_index) - - ordered_responses.extend( - matched_approval_responses[index] - for index in range(len(approval_requests)) - if index not in ordered_approval_indexes - ) - matched_response_ids = {id(response) for response in ordered_responses} | matched_approval_response_ids - filtered_messages: list[Message] = [] - for message in messages: - message.contents = [content for content in message.contents if id(content) not in matched_response_ids] - if message.contents: - filtered_messages.append(message) - filtered_messages.append(Message(role="user", contents=ordered_responses)) - messages[:] = filtered_messages - return False, matched_host_result_ids + filtered_messages: list[Message] = [] + for message in messages: + message.contents = [content for content in message.contents if id(content) not in matched_response_ids] + if message.contents: + filtered_messages.append(message) + filtered_messages.append(Message(role="user", contents=ordered_responses)) + messages[:] = filtered_messages + return False, host_result_ids + return False, set() def _collect_approval_responses( @@ -4018,6 +4016,7 @@ async def _resolve_approval_responses( max_errors: int, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, middleware_pipeline: FunctionMiddlewarePipeline | None = None, settle_dangling_calls: Callable[[Sequence[Content]], Awaitable[None]] | None = None, ) -> _FunctionProcessingResult: @@ -4031,11 +4030,12 @@ async def _resolve_approval_responses( from ._middleware import MiddlewareFailure from ._types import Message + approval_session = invocation_session if approval_session_is_authoritative else None completed_mixed_batch = False - if _has_authoritative_approval_session(invocation_session): + if _has_authoritative_approval_session(approval_session): incomplete_mixed_batch, completed_mixed_batch, host_result_ids = _stage_pending_mixed_pause_responses( prepared_messages, - invocation_session, + approval_session, ) if incomplete_mixed_batch: return _FunctionProcessingResult(errors_in_a_row=errors_in_a_row, action="return") @@ -4047,13 +4047,13 @@ async def _resolve_approval_responses( ) active_pending_ids = ( - set(_load_pending_approval_requests(invocation_session)) - if _has_authoritative_approval_session(invocation_session) + set(_load_pending_approval_requests(approval_session)) + if _has_authoritative_approval_session(approval_session) else None ) - _bind_approval_responses_to_pending_requests(prepared_messages, invocation_session) + _bind_approval_responses_to_pending_requests(prepared_messages, approval_session) if completed_mixed_batch: - state = _get_tool_approval_state(invocation_session, create=False) + state = _get_tool_approval_state(approval_session, create=False) if state is not None: state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) @@ -4066,7 +4066,7 @@ async def _resolve_approval_responses( } if already_approved_responses := _pop_already_approved_approval_responses( - invocation_session, + approval_session, explicit_approval_response_ids, ): prepared_messages.append(Message(role="user", contents=already_approved_responses)) @@ -4151,7 +4151,7 @@ async def _resolve_approval_responses( terminal_contents.extend( request for request_id, request in pending_by_id.items() if request_id not in surfaced_request_ids ) - _store_pending_approval_requests(invocation_session, list(pending_by_id.values())) + _store_pending_approval_requests(approval_session, list(pending_by_id.values())) # 5. Return role-correct output and tell the outer loop whether to return, stop tools, or call the model. executed_function_count = len(execution_result_groups) @@ -4182,8 +4182,10 @@ async def _process_model_function_calls( max_errors: int, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, ) -> _FunctionProcessingResult: """Execute function calls from a newly completed model response.""" + approval_session = invocation_session if approval_session_is_authoritative else None approval_requests = [ content for message in response.messages @@ -4195,7 +4197,7 @@ async def _process_model_function_calls( function_calls = _extract_function_calls(response) if not (function_calls and tools): if approval_requests: - _store_pending_approval_requests(invocation_session, approval_requests) + _store_pending_approval_requests(approval_session, approval_requests) return _FunctionProcessingResult(errors_in_a_row=errors_in_a_row, action="return") # 2. Execute the batch once while preserving each call's result group. @@ -4224,7 +4226,7 @@ async def _process_model_function_calls( if content.type == "function_approval_request" ] if returned_approval_requests: - _store_pending_approval_requests(invocation_session, returned_approval_requests) + _store_pending_approval_requests(approval_session, returned_approval_requests) return processing_result @@ -4385,6 +4387,7 @@ async def _get_response_with_function_invocation( tokenizer: TokenizerProtocol | None, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None, + approval_session_is_authoritative: bool, budget_state: dict[str, Any], max_errors: int, middleware_pipeline: FunctionMiddlewarePipeline | None = None, @@ -4434,6 +4437,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=middleware_pipeline, settle_dangling_calls=settle_approval_replay_calls, ) @@ -4507,6 +4511,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, ) except MiddlewareFailure: # Fail-closed abort: before propagating, settle the batch's calls on a @@ -4598,6 +4603,7 @@ async def _stream_response_with_function_invocation( tokenizer: TokenizerProtocol | None, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None, + approval_session_is_authoritative: bool, budget_state: dict[str, Any], max_errors: int, middleware_pipeline: FunctionMiddlewarePipeline | None = None, @@ -4644,6 +4650,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=middleware_pipeline, settle_dangling_calls=settle_approval_replay_calls, ) @@ -4777,6 +4784,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, ) except MiddlewareFailure: # See the non-streaming loop: settle a service-managed conversation's @@ -4969,9 +4977,12 @@ def get_response( raw_session = request_kwargs.get("session") invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None + approval_session_is_authoritative = ( + request_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, True) is not False + ) if invocation_session is None and requires_session_state: invocation_session = _AgentSession() - setattr(invocation_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, True) + approval_session_is_authoritative = False # Bind one executor with the run's custom arguments, middleware, configuration, and session. execute_function_calls = partial( @@ -4979,6 +4990,7 @@ def get_response( custom_args=additional_function_arguments, config=self.function_invocation_configuration, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=function_middleware_pipeline, host_payload_budget=host_payload_budget, ) @@ -5018,6 +5030,7 @@ def get_response( tokenizer=tokenizer, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, budget_state=budget_state, max_errors=max_errors, middleware_pipeline=function_middleware_pipeline, @@ -5041,6 +5054,7 @@ def finalize_stream(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: tokenizer=tokenizer, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, budget_state=budget_state, max_errors=max_errors, middleware_pipeline=function_middleware_pipeline, diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 0fcd178bb2a..9dea58bf4c2 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3463,6 +3463,243 @@ def approval_func(value: str) -> str: assert approval_arguments == ["expected"] +@pytest.mark.parametrize("include_approval_response", [False, True], ids=["zero-responses", "approval-only"]) +async def test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution( + chat_client_base: SupportsChatGetResponse, + include_approval_response: bool, +) -> None: + """A split stateless mixed batch cannot execute until every response arrives.""" + from agent_framework import FunctionTool + + calls = 0 + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal calls + calls += 1 + return "approved" + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + host_func = FunctionTool(name="host_func", func=None, description="Handled by the caller") + replay_contents = ( + [approval_request.to_function_approval_response(approved=True)] + if include_approval_response + else [Content.from_text("unrelated follow-up")] + ) + messages = [ + Message(role="assistant", contents=[approval_request]), + Message(role="tool", contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message(role="user", contents=replay_contents), + ] + + with pytest.raises( + RuntimeError, + match="A mixed function-call batch requires responses for every approval and Host-owned request", + ): + await chat_client_base.get_response( + messages, + options={"tools": [approval_func, host_func]}, + ) + + assert calls == 0 + + +@pytest.mark.parametrize("metadata_role", ["tool", "user"]) +def test_stateless_mixed_batch_across_message_roles_requires_complete_responses(metadata_role: str) -> None: + """Intervening message roles do not split a mixed batch or change response order.""" + from agent_framework._tools import _stateless_mixed_pause_batch_status + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + approval_response = approval_request.to_function_approval_response(approved=True) + partial_messages = [ + Message(role="assistant", contents=[approval_request]), + Message(role=metadata_role, contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message(role="user", contents=[approval_response]), + ] + + incomplete, _ = _stateless_mixed_pause_batch_status(partial_messages) + + assert incomplete is True + + host_result = Content.from_function_result(call_id="host", result="host result") + host_result.id = "host-occurrence" + complete_messages = [ + Message(role="assistant", contents=[approval_request]), + Message(role=metadata_role, contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message(role="user", contents=[host_result, approval_response]), + ] + + incomplete, host_result_ids = _stateless_mixed_pause_batch_status(complete_messages) + + assert incomplete is False + assert [content.type for content in complete_messages[-1].contents] == [ + "function_approval_response", + "function_result", + ] + assert host_result_ids == {id(complete_messages[-1].contents[1])} + + +async def test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch( + chat_client_base: SupportsChatGetResponse, +) -> None: + """A newer standalone request cannot let an older mixed approval execute early.""" + from agent_framework import FunctionTool + + calls = 0 + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal calls + calls += 1 + return "approved" + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + later_call = Content.from_function_call( + call_id="later", + name="approval_func", + arguments={}, + id="later-occurrence", + ) + later_request = Content.from_function_approval_request( + id="later-occurrence", + function_call=later_call, + ) + host_func = FunctionTool(name="host_func", func=None, description="Handled by the caller") + messages = [ + Message(role="user", contents=["run mixed"]), + Message(role="assistant", contents=[approval_request, host_request]), + Message(role="user", contents=[approval_request.to_function_approval_response(approved=True)]), + Message(role="assistant", contents=[later_request]), + Message(role="user", contents=["unrelated follow-up"]), + ] + + with pytest.raises( + RuntimeError, + match="A mixed function-call batch requires responses for every approval and Host-owned request", + ): + await chat_client_base.get_response( + messages, + options={"tools": [approval_func, host_func]}, + ) + + assert calls == 0 + + +async def test_completed_split_stateless_mixed_batch_is_inert_on_later_turn( + chat_client_base: SupportsChatGetResponse, +) -> None: + """A terminal result prevents completed split mixed approvals from replaying.""" + from agent_framework import FunctionTool + + calls = 0 + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal calls + calls += 1 + return "approved" + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + host_result = Content.from_function_result(call_id="host", result="host result") + host_result.id = "host-occurrence" + approval_result = Content.from_function_result(call_id="approval", result="approved") + host_func = FunctionTool(name="host_func", func=None, description="Handled by the caller") + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse(messages=Message(role="assistant", contents=["later response"])), + ] + messages = [ + Message(role="user", contents=["run mixed"]), + Message(role="assistant", contents=[approval_request]), + Message(role="tool", contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message( + role="user", + contents=[ + approval_request.to_function_approval_response(approved=True), + host_result, + ], + ), + Message(role="tool", contents=[approval_result]), + Message(role="assistant", contents=["done"]), + Message(role="user", contents=["later"]), + ] + + response = await chat_client_base.get_response( + messages, + options={"tools": [approval_func, host_func]}, + ) + + assert response.text == "later response" + assert calls == 0 + + def test_active_mixed_pause_ignores_historical_host_requests() -> None: """Only the session-recorded mixed batch participates in response correlation.""" from agent_framework._tools import ( diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 95f65e5a467..b3b88359215 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -24,6 +24,7 @@ ChatResponse, ChatResponseUpdate, Content, + ContextProvider, FileHistoryProvider, FunctionInvocationContext, FunctionMiddleware, @@ -31,6 +32,8 @@ InMemoryHistoryProvider, Message, ResponseStream, + SessionContext, + SupportsAgentRun, ToolApprovalMiddleware, ToolApprovalState, create_always_approve_tool_response, @@ -134,6 +137,136 @@ def safe_tool() -> str: ] +async def test_sessionless_approval_resume_with_context_provider_uses_message_authority( + chat_client_base: MockBaseChatClient, +) -> None: + """A framework-created context-provider session must not become approval authority.""" + calls: list[str] = [] + + class NoOpContextProvider(ContextProvider): + pass + + @tool(name="approved_tool", approval_mode="always_require") + def approved_tool() -> str: + calls.append("approved") + return "approved" + + agent = Agent( + client=chat_client_base, + tools=[approved_tool], + context_providers=[NoOpContextProvider("context")], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="approved-call", + name="approved_tool", + arguments="{}", + id="approved-occurrence", + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + first = await agent.run("request approval") + approval_request = first.user_input_requests[0] + resumed = await agent.run([ + Message(role="user", contents=["request approval"]), + *first.messages, + Message(role="user", contents=[approval_request.to_function_approval_response(True)]), + ]) + + assert calls == ["approved"] + assert [(message.role, [content.type for content in message.contents]) for message in resumed.messages] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + + +async def test_framework_created_session_becomes_authoritative_when_reused_by_caller( + chat_client_base: MockBaseChatClient, +) -> None: + """A retained framework session must fail closed when later supplied explicitly.""" + calls: list[int] = [] + retained_sessions: list[AgentSession] = [] + + class RetainingContextProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, context, state + retained_sessions.append(session) + + @tool(name="approved_tool", approval_mode="always_require") + def approved_tool(value: int) -> str: + calls.append(value) + return "approved" + + agent = Agent( + client=chat_client_base, + tools=[approved_tool], + context_providers=[RetainingContextProvider("context")], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="approved-call", + name="approved_tool", + arguments={"value": 1}, + id="approved-occurrence", + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + first = await agent.run("request approval") + assert calls == [] + assert "tool_approval" not in retained_sessions[0].state + approval_request = first.user_input_requests[0] + original_call = _function_call(approval_request) + assert original_call.call_id is not None + assert original_call.name is not None + assert approval_request.id is not None + tampered_call = Content.from_function_call( + call_id=original_call.call_id, + name=original_call.name, + arguments={"value": 999}, + id=original_call.id, + ) + tampered_response = Content.from_function_approval_response( + approved=True, + id=approval_request.id, + function_call=tampered_call, + ) + + await agent.run( + [ + Message(role="user", contents=["request approval"]), + *first.messages, + Message(role="user", contents=[tampered_response]), + ], + session=retained_sessions[0], + ) + + assert calls == [] + assert retained_sessions[1] is retained_sessions[0] + + class _StringApprovalValue(str, Enum): ALPHA = "alpha" @@ -2116,6 +2249,53 @@ def second_guarded() -> str: assert state["queued_approval_requests"] == [] +@pytest.mark.parametrize("approval_first", [True, False], ids=["approval-first", "host-first"]) +async def test_tool_approval_middleware_preserves_split_streamed_mixed_batch_order( + chat_client_base: MockBaseChatClient, + approval_first: bool, +) -> None: + """Separately streamed approval and Host pauses remain one ordered batch.""" + from agent_framework import FunctionTool + from agent_framework._middleware import AgentContext + + @tool(name="guarded", approval_mode="always_require") + def guarded() -> str: + return "guarded" + + host_tool = FunctionTool(name="host_read", func=None, description="Handled by the caller") + agent = Agent(client=chat_client_base, tools=[guarded, host_tool]) + approval_call = Content.from_function_call(call_id="guarded", name="guarded", arguments={}) + approval_request = Content.from_function_approval_request(id="guarded", function_call=approval_call) + host_request = Content.from_function_call(call_id="host", name="host_read", arguments={}) + host_request.user_input_request = True + ordered_requests = [approval_request, host_request] if approval_first else [host_request, approval_request] + context = AgentContext( + agent=agent, + messages=[], + session=AgentSession(), + stream=True, + ) + middleware = ToolApprovalMiddleware() + state = ToolApprovalState() + + async def inner_stream(): + for request in ordered_requests: + yield AgentResponseUpdate(role="assistant", contents=[request]) + + async def call_next() -> None: + context.result = ResponseStream(inner_stream(), finalizer=AgentResponse.from_updates) + + response_stream = middleware._process_stream( # pyright: ignore[reportPrivateUsage] + context, + call_next, + state, + ) + updates = [update async for update in response_stream] + + assert [content for update in updates for content in update.user_input_requests] == ordered_requests + assert state.queued_approval_requests == [] + + @pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) async def test_tool_approval_middleware_policy_approval_reclassifies_host_tool( chat_client_base: MockBaseChatClient, From 3b110ade39ea798486b1060f40ee9e2c6329d0ae Mon Sep 17 00:00:00 2001 From: lsmlhi_25 Date: Thu, 17 Sep 2026 08:22:14 +0000 Subject: [PATCH 03/26] Python: include checkpoint_id on AG-UI interrupt metadata (#8150) (#8163) * Python: include checkpoint_id on AG-UI interrupt metadata (#8150) Attach the pause workflow checkpoint id to RUN_FINISHED interrupt metadata.agent_framework so multi-worker clients can resume via forwardedProps.checkpoint_id without a side channel. * Python: Prefer runner pause checkpoint over shared get_latest (#8150) * Python: Fix ag-ui lint and ty ignores for #8150 tests * Python: Core resolve_pause_checkpoint_id for AG-UI (#8150) Move pause-checkpoint selection into Workflow and require a run-scoped baseline so leftover runner ids are not advertised across runs. * fix(python): resolve pyright failures in pause checkpoint lookup Unblocks Package Checks / merge-gatekeeper on PR #8163. * fix(ag-ui): address #8163 review on pause checkpoint ownership (#2) - Workflow captures run baseline / restored id inside run(); resolve uses them so callers need not thread baseline_checkpoint_id. - Exclude restored checkpoint from pause candidates after resume. - Drop issue reference in docstring; emit RUN_FINISHED via existing _build_run_finished_event after attaching pause id. - Move core pause-resolve coverage into test_workflow.py (no new test file). * fix(ag-ui): resume builder-storage pause IDs without AG-UI storage When checkpoint_id+resume is supplied, load pending requests from the workflow's effective builder/runtime storage if AG-UI was not given an explicit checkpoint_storage, so emitted pause IDs remain round-trippable. * fix(ag-ui): silence ty invalid-argument-type on builder-storage resume test Match existing yield_output ignore markers so Test Typing Checks accept the new resume round-trip coverage. * fix(ag-ui): resume builder storage via AgentFrameworkWorkflow.run() Allow forwardedProps.checkpoint_id when only WorkflowBuilder storage is configured, so wrapper/endpoint hosts round-trip pause ids without duplicating checkpoint_storage. --------- Co-authored-by: minelhi <3417378192@qq.com> Co-authored-by: LI <2484593937@qq.com> Co-authored-by: Eduard van Valkenburg --- .../ag-ui/agent_framework_ag_ui/_workflow.py | 29 ++- .../agent_framework_ag_ui/_workflow_run.py | 180 +++++++++++++- .../ag-ui/tests/ag_ui/test_workflow_agent.py | 35 ++- .../ag-ui/tests/ag_ui/test_workflow_run.py | 225 ++++++++++++++++++ .../agent_framework/_workflows/_workflow.py | 99 ++++++++ .../core/tests/workflow/test_workflow.py | 106 +++++++++ 6 files changed, 655 insertions(+), 19 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 539ddd21855..33d3bede30e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -525,6 +525,12 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: (``forwarded_props: {"checkpoint_id": ...}``), which restores the persisted workflow state instead of starting a fresh turn. + Resume also works when storage lives only on the resolved workflow + (``WorkflowBuilder(checkpoint_storage=...)``) and this wrapper / endpoint omits + ``checkpoint_storage``, so emitted pause ids remain round-trippable for hosts that + expose the workflow through ``AgentFrameworkWorkflow`` or + ``add_agent_framework_fastapi_endpoint`` without duplicating storage configuration. + Note: Checkpointing (the ``agent_framework`` workflow checkpoint mechanism) is independent from AG-UI Thread Snapshot persistence (``snapshot_store``). @@ -547,11 +553,6 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: checkpoint_storage = self.checkpoint_storage checkpoint_id = _checkpoint_id_from_input(input_data) - if checkpoint_id is not None and checkpoint_storage is None: - raise ValueError( - "Resuming from a checkpoint requires checkpoint_storage to be configured on " - "AgentFrameworkWorkflow (or the AG-UI endpoint)." - ) supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") request_owner = (snapshot_scope, str(supplied_thread_id) if supplied_thread_id is not None else None) @@ -561,6 +562,15 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: if interrupt.get("id") is not None } workflow = self._resolve_workflow(thread_id, snapshot_scope) + # Prefer wrapper/endpoint storage; otherwise allow builder/runtime storage on the + # resolved workflow so pause ids emitted without AG-UI storage still resume. + if checkpoint_id is not None and checkpoint_storage is None: + if not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] + raise ValueError( + "Resuming from a checkpoint requires checkpoint_storage to be configured on " + "AgentFrameworkWorkflow (or the AG-UI endpoint), or WorkflowBuilder " + "checkpoint storage on the workflow instance." + ) live_pending_events = await _pending_request_events(self.workflow) if self.workflow is not None else {} if self.workflow is not None and checkpoint_id is None: for request_event in live_pending_events.values(): @@ -572,9 +582,14 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: code="WORKFLOW_RESUME_NOT_FOUND", ) return - if checkpoint_id is not None and checkpoint_storage is not None: + if checkpoint_id is not None: try: - checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + else: + checkpoint = await workflow._runner.context.load_checkpoint(checkpoint_id) # pyright: ignore[reportPrivateUsage] + if checkpoint is None: + raise LookupError(f"checkpoint '{checkpoint_id}' was not found") except Exception as exc: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield RunErrorEvent( diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 9ee70075f2b..9a5126cc813 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -8,7 +8,7 @@ import json import logging import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from functools import partial from types import UnionType from typing import Any, Union, cast, get_args, get_origin, get_type_hints @@ -61,6 +61,8 @@ logger = logging.getLogger(__name__) +_BASELINE_OMITTED = object() + _PUBLIC_WORKFLOW_ERROR_MESSAGE = "Workflow execution failed." @@ -130,7 +132,10 @@ def _workflow_interrupt_value(request_data: Any) -> Any: return {"data": safe_request_data} -def _workflow_interrupt_metadata(request_payload: dict[str, Any], value: Any) -> dict[str, Any]: +def _workflow_interrupt_metadata( + request_payload: dict[str, Any], + value: Any, +) -> dict[str, Any]: """Build Agent Framework metadata for workflow request_info interrupts.""" agent_framework_metadata = { key: make_json_safe(value) @@ -148,6 +153,99 @@ def _workflow_interrupt_metadata(request_payload: dict[str, Any], value: Any) -> return {"agent_framework": agent_framework_metadata} +def _attach_checkpoint_id_to_interrupts( + interrupts: list[dict[str, Any]], + checkpoint_id: str | None, +) -> list[dict[str, Any]]: + """Attach ``checkpoint_id`` to each interrupt's ``metadata.agent_framework``. + + Multi-worker hosts need the pause checkpoint on the wire so the next resume can pass + ``forwardedProps.checkpoint_id`` without a side-channel lookup. No-op when checkpointing + is inactive or the id is already present. + """ + if not checkpoint_id or not interrupts: + return interrupts + + attached: list[dict[str, Any]] = [] + for interrupt in interrupts: + entry = dict(interrupt) + metadata = entry.get("metadata") + if isinstance(metadata, dict): + metadata = dict(metadata) + else: + metadata = {} + agent_framework = metadata.get("agent_framework") + if isinstance(agent_framework, dict): + agent_framework = dict(agent_framework) + else: + agent_framework = {} + agent_framework.setdefault("checkpoint_id", checkpoint_id) + metadata["agent_framework"] = agent_framework + entry["metadata"] = metadata + attached.append(entry) + return attached + + +def _interrupt_request_ids(interrupts: list[dict[str, Any]]) -> set[str]: + return {str(item["id"]) for item in interrupts if item.get("id") is not None} + + +async def _pause_checkpoint_id_for_interrupts( + *, + workflow: Workflow, + checkpoint_storage: CheckpointStorage | None, + interrupts: list[dict[str, Any]], + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: Any = _BASELINE_OMITTED, +) -> str | None: + """Resolve the pause checkpoint for *this* run's interrupts via core. + + When ``baseline_checkpoint_id`` is omitted, core uses the baseline captured at + ``workflow.run()`` start. Pass ``None`` explicitly for short-circuit paths that + did not call ``run()`` and should advertise the current runner id when present. + """ + if not interrupts: + return None + + resolve = getattr(workflow, "resolve_pause_checkpoint_id", None) + if not callable(resolve): + return None + + # getattr returns a plain object to the type checker; cast to an awaitable callable. + resolve_fn = cast( + Callable[..., Awaitable[str | None]], + resolve, + ) + kwargs: dict[str, Any] = { + "checkpoint_storage": checkpoint_storage, + "known_checkpoint_id": known_checkpoint_id, + } + if baseline_checkpoint_id is not _BASELINE_OMITTED: + kwargs["baseline_checkpoint_id"] = baseline_checkpoint_id + return await resolve_fn(_interrupt_request_ids(interrupts), **kwargs) + + +async def _interrupts_with_pause_checkpoint( + *, + interrupts: list[dict[str, Any]], + workflow: Workflow, + checkpoint_storage: CheckpointStorage | None, + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: Any = _BASELINE_OMITTED, +) -> list[dict[str, Any]]: + """Attach a run-scoped pause checkpoint id to interrupts when available.""" + if not interrupts: + return interrupts + pause_checkpoint_id = await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=checkpoint_storage, + interrupts=interrupts, + known_checkpoint_id=known_checkpoint_id, + baseline_checkpoint_id=baseline_checkpoint_id, + ) + return _attach_checkpoint_id_to_interrupts(interrupts, pause_checkpoint_id) + + async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: """Best-effort retrieval of pending request_info events from workflow context.""" runner_context = getattr(workflow, "_runner_context", None) @@ -171,7 +269,9 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: async def _pending_request_events_from_checkpoint( checkpoint_id: str, - checkpoint_storage: CheckpointStorage, + checkpoint_storage: CheckpointStorage | None = None, + *, + workflow: Any | None = None, ) -> dict[str, Any]: """Read pending request_info events from a persisted checkpoint without restoring it. @@ -182,16 +282,34 @@ async def _pending_request_events_from_checkpoint( exposes them without running any executor ``on_checkpoint_restore`` hook; the single ``workflow.run(checkpoint_id=...)`` then performs the one real restore, so the restore -- and every custom restore hook -- runs exactly once per resume. + + ``checkpoint_storage`` is preferred when provided. Otherwise the workflow's + effective builder/runtime storage is used when ``has_checkpointing()`` is true, + so AG-UI can round-trip pause IDs emitted from ``WorkflowBuilder(checkpoint_storage=...)`` + without requiring a duplicate AG-UI storage argument. """ try: - checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + else: + context = getattr(getattr(workflow, "_runner", None), "context", None) + if context is None or not context.has_checkpointing(): + raise ValueError( + "Resuming a checkpoint with an AG-UI resume payload requires checkpoint_storage " + "(or WorkflowBuilder checkpoint storage on the workflow instance)." + ) + checkpoint = await context.load_checkpoint(checkpoint_id) + except ValueError: + raise except Exception: logger.warning( "Could not load checkpoint for resume-response coercion; the core run will surface any error.", exc_info=True, ) return {} - return dict(checkpoint.pending_request_info_events) + if checkpoint is None: + return {} + return dict(checkpoint.pending_request_info_events or {}) def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None: @@ -1062,9 +1180,15 @@ async def run_workflow_stream( # pure checkpoint restore still surfaces its pending interrupts instead of tripping # the "resume required" contract. if checkpoint_id is not None and resume_payload is not None: - if checkpoint_storage is None: + # Prefer the explicit AG-UI storage argument; otherwise allow the workflow's + # builder/runtime storage so builder-emitted pause IDs remain round-trippable. + if checkpoint_storage is None and not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] raise ValueError("Resuming a checkpoint with an AG-UI resume payload requires checkpoint_storage.") - pending_before_run = await _pending_request_events_from_checkpoint(checkpoint_id, checkpoint_storage) + pending_before_run = await _pending_request_events_from_checkpoint( + checkpoint_id, + checkpoint_storage, + workflow=workflow, + ) else: pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) @@ -1140,12 +1264,30 @@ async def run_workflow_stream( interrupt_event_value = _workflow_interrupt_event_value(request_payload) if interrupt_event_value is not None: yield CustomEvent(name=_INTERRUPT_CARD_EVENT_NAME, value=interrupt_event_value) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=pending_interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + baseline_checkpoint_id=None, + ), + ) return if checkpoint_id is None and not responses and not messages: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=pending_interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + baseline_checkpoint_id=None, + ), + ) return def _drain_open_message() -> list[TextMessageEndEvent]: @@ -1281,7 +1423,15 @@ def _drain_open_blocks() -> list[BaseEvent]: yield end_event if not interrupts: interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow))) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + ), + ) terminal_emitted = True elif state_value not in _TERMINAL_STATES: yield CustomEvent(name="status", value={"state": state_value}) @@ -1447,4 +1597,12 @@ def _drain_open_blocks() -> list[BaseEvent]: if not terminal_emitted and not run_error_emitted: if not interrupts: interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow))) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + ), + ) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 8702adcd047..615d3ce4f1a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -260,7 +260,7 @@ async def test_workflow_run_reads_checkpoint_id_from_camelcase_forwarded_props() async def test_workflow_resume_without_checkpoint_storage_raises() -> None: - """Requesting a checkpoint resume without configured storage should fail loudly.""" + """Requesting a checkpoint resume without any storage should fail loudly.""" workflow = _build_multi_superstep_workflow() agent = AgentFrameworkWorkflow(workflow=workflow) @@ -275,6 +275,39 @@ async def test_workflow_resume_without_checkpoint_storage_raises() -> None: ) +async def test_workflow_wrapper_resumes_builder_storage_without_agui_storage() -> None: + """Builder-owned storage must round-trip through AgentFrameworkWorkflow.run() without wrapper storage.""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow(storage) + # Host configures storage only on the builder; AG-UI wrapper / endpoint omit it. + agent = AgentFrameworkWorkflow(workflow=workflow) + + first_events = await _run( + agent, + {"thread_id": "thread-builder-cp", "messages": [{"role": "user", "content": "start"}]}, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the builder-storage run to create at least one checkpoint" + resume_checkpoint_id = checkpoints[0].checkpoint_id + + resume_events = await _run( + agent, + { + "thread_id": "thread-builder-cp", + "messages": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, + }, + ) + resumed_types = [event.type for event in resume_events] + assert "RUN_FINISHED" in resumed_types + assert "RUN_ERROR" not in resumed_types + + async def test_workflow_run_without_checkpointing_is_unchanged() -> None: """Existing run(input_data) calls keep working unchanged when no checkpoint args are given.""" workflow = _build_multi_superstep_workflow() diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index d9cb44fe01c..fc79a1ae07f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -98,6 +98,229 @@ def _interrupt_metadata_value(interrupt: dict[str, Any]) -> dict[str, Any]: return cast(dict[str, Any], value) +def test_attach_checkpoint_id_to_interrupts_setdefault() -> None: + """Issue #8150: attach pause checkpoint_id without overwriting an existing one.""" + from agent_framework_ag_ui._workflow_run import _attach_checkpoint_id_to_interrupts + + empty = _attach_checkpoint_id_to_interrupts([{"id": "r1"}], None) + assert empty == [{"id": "r1"}] + + attached = _attach_checkpoint_id_to_interrupts( + [{"id": "r1", "metadata": {"agent_framework": {"type": "workflow_request_info"}}}], + "cp-123", + ) + assert attached[0]["metadata"]["agent_framework"]["checkpoint_id"] == "cp-123" + assert attached[0]["metadata"]["agent_framework"]["type"] == "workflow_request_info" + + preserved = _attach_checkpoint_id_to_interrupts(attached, "cp-other") + assert preserved[0]["metadata"]["agent_framework"]["checkpoint_id"] == "cp-123" + + +@pytest.mark.asyncio +async def test_pause_checkpoint_id_ignores_competing_shared_latest() -> None: + """Prefer this runner's pause checkpoint over a newer shared get_latest() winner.""" + from agent_framework import WorkflowCheckpoint + + from agent_framework_ag_ui._run_common import _build_run_finished_event + from agent_framework_ag_ui._workflow_run import ( + _interrupts_with_pause_checkpoint, + _pause_checkpoint_id_for_interrupts, + ) + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345"}, + ) + approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call) + await ctx.request_info(approval_request, Content, request_id="approval-1") + + @response_handler + async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None: + del original_request, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + first_events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "go"}]}, + workflow, + checkpoint_storage=storage, + ) + ] + first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0] + interrupt_payload = _interrupts_from_run_finished(first_finished) + pause_id = interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] + + # Inject a newer shared checkpoint that does not cover this interrupt (other owner / stale). + competing = WorkflowCheckpoint( + workflow_name=workflow.name, + graph_signature_hash="competing", + pending_request_info_events={}, + timestamp="9999-01-01T00:00:00+00:00", + ) + await storage.save(competing) + latest = await storage.get_latest(workflow_name=workflow.name) + assert latest is not None + assert latest.checkpoint_id == competing.checkpoint_id + + # Workflow-owned baseline from the pause run should already allow advertising pause_id. + resolved = await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=storage, + interrupts=interrupt_payload, + ) + assert resolved == pause_id + + rebuilt = _build_run_finished_event( + "run-1", + "thread-1", + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupt_payload, + workflow=workflow, + checkpoint_storage=storage, + ), + ) + rebuilt_interrupts = _interrupts_from_run_finished(rebuilt) + assert rebuilt_interrupts[0]["metadata"]["agent_framework"]["checkpoint_id"] == pause_id + + +@pytest.mark.asyncio +async def test_builder_checkpoint_storage_attaches_id_without_run_arg() -> None: + """WorkflowBuilder(checkpoint_storage=...) alone must still advertise pause checkpoint_id.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345"}, + ) + approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call) + await ctx.request_info(approval_request, Content, request_id="approval-1") + + @response_handler + async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None: + del original_request, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + # Deliberately omit checkpoint_storage= on the AG-UI entrypoint (builder path only). + events = [ + event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow) + ] + finished = [event for event in events if event.type == "RUN_FINISHED"][0] + interrupt_payload = _interrupts_from_run_finished(finished) + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + assert checkpoints + assert interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] == checkpoints[-1].checkpoint_id + + +@pytest.mark.asyncio +async def test_builder_checkpoint_storage_resume_round_trips_without_agui_storage() -> None: + """Pause IDs from builder storage must resume with resume payload even without AG-UI storage.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + await ctx.request_info("need-input", str, request_id="req-1") + + @response_handler + async def handle(self, original_request: str, response: str, ctx: WorkflowContext) -> None: + del original_request + await ctx.yield_output(f"got:{response}") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + + pause_events = [ + event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow) + ] + finished = [event for event in pause_events if event.type == "RUN_FINISHED"][0] + pause_id = _interrupts_from_run_finished(finished)[0]["metadata"]["agent_framework"]["checkpoint_id"] + + # Cold resume: omit AG-UI checkpoint_storage; rely on builder storage only. + resume_events = [ + event + async for event in run_workflow_stream( + { + "messages": [], + "resume": {"interrupts": [{"id": "req-1", "value": "ok"}]}, + "forwarded_props": {"checkpoint_id": pause_id}, + }, + workflow, + ) + ] + assert "RUN_ERROR" not in [event.type for event in resume_events] + assert "RUN_FINISHED" in [event.type for event in resume_events] + text = "".join(getattr(event, "delta", "") for event in resume_events if event.type == "TEXT_MESSAGE_CONTENT") + assert "got:ok" in text + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_is_run_scoped_without_storage() -> None: + """Stale runner ids must not be advertised when baseline shows this run did not persist.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + await ctx.request_info("need-input", str, request_id="req-1") + + @response_handler + async def handle(self, original_request: str, response: str, ctx: WorkflowContext) -> None: + del original_request, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + # Simulate a leftover id from a prior run on a workflow with no checkpoint storage. + workflow._runner._previous_checkpoint_id = "stale-from-prior-run" # pyright: ignore[reportPrivateUsage] + + from agent_framework_ag_ui._workflow_run import _pause_checkpoint_id_for_interrupts + + interrupts = [{"id": "req-1", "value": "need-input"}] + workflow._run_baseline_checkpoint_id = "stale-from-prior-run" # pyright: ignore[reportPrivateUsage] + assert ( + await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=None, + interrupts=interrupts, + ) + is None + ) + assert ( + await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=None, + interrupts=interrupts, + baseline_checkpoint_id=None, + ) + == "stale-from-prior-run" + ) + + async def test_workflow_run_maps_custom_and_text_events(): """Custom workflow events and yielded text are mapped to AG-UI events.""" @@ -843,6 +1066,8 @@ async def handle_approval(self, original_request: Content, response: Content, ct ) assert checkpoints, "expected the interrupted run to create a checkpoint" resume_checkpoint_id = checkpoints[-1].checkpoint_id + # Issue #8150: interrupt metadata must carry the pause checkpoint for multi-worker resume. + assert interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] == resume_checkpoint_id # Resume on a FRESH workflow instance so no pending requests exist in memory until # the checkpoint is restored -- a cold restore, as after a process restart. diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index afbf45af11b..a24a2fade38 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -410,6 +410,12 @@ def __init__( # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). + # Captured at the start of each ``_run_core`` so ``resolve_pause_checkpoint_id`` + # can tell a newly persisted pause from a leftover / restored id. + self._run_baseline_checkpoint_id: str | None = None + self._restored_checkpoint_id: str | None = None + @property def status(self) -> WorkflowRunState: """Return the current run-level status of this workflow instance. @@ -907,6 +913,12 @@ async def _run_core( if checkpoint_storage is not None: self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + # Own the run boundary for pause-checkpoint resolution: capture the runner id + # before this run advances it, and remember an incoming restore so it is not + # re-advertised if a later pause save fails. + self._run_baseline_checkpoint_id = self.get_last_checkpoint_id() + self._restored_checkpoint_id = str(checkpoint_id) if checkpoint_id is not None else None + try: # Async validation: a fresh-message run is only allowed when the # runner context has fully drained from any prior run. If it still @@ -1308,6 +1320,93 @@ def output_types(self) -> list[type[Any] | types.UnionType]: return list(output_types) + def get_last_checkpoint_id(self) -> str | None: + """Return the checkpoint id last persisted or restored by this workflow runner.""" + checkpoint_id = self._runner._previous_checkpoint_id # pyright: ignore[reportPrivateUsage] + return str(checkpoint_id) if checkpoint_id is not None else None + + async def resolve_pause_checkpoint_id( + self, + request_ids: Collection[str], + *, + checkpoint_storage: CheckpointStorage | None = None, + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: str | None | object = _MISSING, + ) -> str | None: + """Resolve the persisted pause checkpoint that covers ``request_ids``. + + Prefers this runner's last-saved id when it advanced past the run baseline + (captured automatically at ``run()`` start), over shared ``get_latest(workflow_name=...)``. + Storage precedence is the run argument, else the runner's effective + (runtime / builder) storage. When storage is available, candidates are accepted + only if their ``pending_request_info_events`` cover ``request_ids``. + + An incoming restored checkpoint id is excluded so a failed pause save after + resume cannot re-advertise pre-response state. + + Args: + request_ids: Pending request_info ids that must be present on the checkpoint. + checkpoint_storage: Optional storage override for this lookup. + known_checkpoint_id: Fallback id (for example a cold-resume short-circuit). + baseline_checkpoint_id: Optional override for the pre-run runner id. When omitted, + uses the baseline captured by the most recent ``run()``. Pass ``None`` + explicitly to treat any current runner id as newly advanced. + + Returns: + A checkpoint id suitable for durable resume, or ``None`` when none is safe. + """ + ids = {str(request_id) for request_id in request_ids if request_id} + if not ids: + return None + + if baseline_checkpoint_id is _MISSING: + baseline_checkpoint_id = self._run_baseline_checkpoint_id + + # Prefer explicit storage; otherwise load via public RunnerContext APIs + # (Protocol has no private `_get_effective_checkpoint_storage`). + storage = checkpoint_storage + use_context_storage = storage is None and self._runner.context.has_checkpointing() + + current = self.get_last_checkpoint_id() + excluded: set[str] = set() + if self._restored_checkpoint_id is not None: + excluded.add(self._restored_checkpoint_id) + + # Run-scoped: do not advertise a pre-run leftover when this run did not persist. + runner_candidate: str | None = None + if current is not None and current != baseline_checkpoint_id and current not in excluded: + runner_candidate = current + + candidates: list[str] = [] + if runner_candidate is not None: + candidates.append(runner_candidate) + if ( + known_checkpoint_id is not None + and known_checkpoint_id not in candidates + and known_checkpoint_id not in excluded + ): + candidates.append(str(known_checkpoint_id)) + + if storage is None and not use_context_storage: + # Without storage we cannot prove coverage; only advertise a run-scoped runner id. + return runner_candidate + + for candidate in candidates: + try: + if storage is not None: + checkpoint = await storage.load(candidate) + else: + checkpoint = await self._runner.context.load_checkpoint(candidate) + except Exception: # pragma: no cover - storage/type drift + logger.debug("Could not load pause checkpoint candidate %s", candidate, exc_info=True) + continue + if checkpoint is None: + continue + pending = checkpoint.pending_request_info_events or {} + if ids.issubset({str(key) for key in dict(pending)}): + return candidate + return None + async def cancel_pending_requests( self, request_ids: Collection[str], diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index bf5487ce0f8..fa96ee9d60f 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -26,6 +26,7 @@ Message, ResponseStream, WorkflowBuilder, + WorkflowCheckpoint, WorkflowCheckpointException, WorkflowContext, WorkflowConvergenceException, @@ -1769,3 +1770,108 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None # endregion + + +# --------------------------------------------------------------------------- +# Pause checkpoint resolution (AG-UI interrupt metadata) +# --------------------------------------------------------------------------- + +class _ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + function_call = Content.from_function_call( + call_id="refund-call", + name="submit_refund", + arguments={"order_id": "12345"}, + ) + approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call) + await ctx.request_info(approval_request, Content, request_id="approval-1") + + @response_handler + async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None: + del original_request, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_prefers_runner_over_shared_latest() -> None: + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + assert pause_id is not None + + competing = WorkflowCheckpoint( + workflow_name=workflow.name, + graph_signature_hash="competing", + pending_request_info_events={}, + timestamp="9999-01-01T00:00:00+00:00", + ) + await storage.save(competing) + + # Workflow owns the run baseline captured at run() start. + resolved = await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + ) + assert resolved == pause_id + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_requires_run_scoped_change_without_storage() -> None: + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor()).build() + workflow._runner._previous_checkpoint_id = "leftover" # pyright: ignore[reportPrivateUsage] + workflow._run_baseline_checkpoint_id = "leftover" # pyright: ignore[reportPrivateUsage] + + assert await workflow.resolve_pause_checkpoint_id({"approval-1"}) is None + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + baseline_checkpoint_id=None, + ) + == "leftover" + ) + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_excludes_restored_checkpoint() -> None: + """After resume, the restored id must not be re-advertised if a new pause save fails.""" + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + + async for _ in workflow.run("go", stream=True): + pass + restored = workflow.get_last_checkpoint_id() + assert restored is not None + + # Simulate a resume that restored `restored` but did not persist a newer pause. + workflow._run_baseline_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + workflow._restored_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + workflow._runner._previous_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + known_checkpoint_id=restored, + ) + is None + ) + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_uses_builder_storage() -> None: + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + resolved = await workflow.resolve_pause_checkpoint_id({"approval-1"}) + assert resolved == pause_id From 52e03cbd01ec9ecfa2b8b40c0c0bacf00f469bb6 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:36:00 +0000 Subject: [PATCH 04/26] .NET/Python: Refine skill frontmatter parsing (#8430) * Reject ambiguous skill frontmatter fields in Python and .NET Validate recognized top-level field casing and uniqueness while preserving existing scalar parsing and nested metadata behavior. Add file and MCP archive regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Keep first duplicate skill metadata value in Python and .NET Warn on duplicate metadata entries without rejecting skills, preserve existing key comparison rules, and expand scalar and archive regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Preserve unset optional skill frontmatter fields Keep bare empty optional fields null while retaining duplicate and casing validation in Python and .NET. Add file and MCP archive regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Support quoted skill frontmatter property names Normalize single- and double-quoted root names before existing field validation in Python and .NET. Support quoted metadata headers and expand scalar and archive regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Use PyYAML for Python skill frontmatter parsing Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com> --- .../Skills/AgentSkillFrontmatter.cs | 4 + .../Skills/File/AgentFileSkillsSource.cs | 86 +- .../AgentMcpSkillsSourceArchiveTests.cs | 180 ++++ .../AgentSkills/FileAgentSkillLoaderTests.cs | 647 ++++++++++++++ python/packages/core/AGENTS.md | 1 + .../packages/core/agent_framework/_skills.py | 301 +++---- python/packages/core/pyproject.toml | 2 + .../core/tests/core/test_mcp_skills.py | 155 ++++ .../packages/core/tests/core/test_skills.py | 803 +++++++++++++++++- python/uv.lock | 4 + 10 files changed, 2015 insertions(+), 168 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs b/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs index 136d4d6833a..06e0e7853da 100644 --- a/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs +++ b/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs @@ -112,6 +112,10 @@ public string? Compatibility /// /// Gets or sets the arbitrary key-value metadata for this skill. /// + /// + /// Keys are compared case-insensitively. When parsed from a SKILL.md file, duplicate + /// entries retain the first value and key spelling and produce warnings without rejecting the skill. + /// public AdditionalPropertiesDictionary? Metadata { get; set; } /// diff --git a/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs b/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs index b22ea74c93e..5484348c59e 100644 --- a/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs +++ b/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs @@ -22,7 +22,10 @@ namespace Microsoft.Agents.AI; /// /// Searches directories recursively (up to 2 levels deep) for SKILL.md files. /// Symbolic links and reparse points below configured roots are not followed during skill discovery. -/// Each file is validated for YAML frontmatter. Resource and script files are discovered by scanning the skill +/// Recognized top-level frontmatter fields must use lowercase names and must not be repeated. +/// Within the optional metadata mapping, keys are compared case-insensitively. The first +/// value is retained for duplicate keys, and subsequent entries produce warnings without rejecting the skill. +/// Resource and script files are discovered by scanning the skill /// directory for files with matching extensions. Invalid resources are skipped with logged warnings. /// Resource and script paths are checked against path traversal and symlink escape attacks. /// Discovered files are revalidated against their trusted skill directory immediately before use. @@ -36,19 +39,31 @@ public sealed partial class AgentFileSkillsSource : AgentSkillsSource private static readonly string[] s_defaultScriptExtensions = [".py", ".js", ".sh", ".ps1", ".cs", ".csx"]; private static readonly string[] s_defaultResourceExtensions = [".md", ".json", ".yaml", ".yml", ".csv", ".xml", ".txt"]; + // Case-insensitive lookup identifies known fields; values preserve the required spelling for validation. + private static readonly Dictionary s_frontmatterFieldNames = new(StringComparer.OrdinalIgnoreCase) + { + ["name"] = "name", + ["description"] = "description", + ["license"] = "license", + ["compatibility"] = "compatibility", + ["metadata"] = "metadata", + ["allowed-tools"] = "allowed-tools", + }; + // Matches YAML frontmatter delimited by "---" lines. Group 1 = content between delimiters. // Multiline makes ^/$ match line boundaries; Singleline makes . match newlines across the block. // The \uFEFF? prefix allows an optional UTF-8 BOM that some editors prepend. private static readonly Regex s_frontmatterRegex = new(@"\A\uFEFF?^---\s*$(.+?)^---\s*$", RegexOptions.Multiline | RegexOptions.Singleline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); - // Matches top-level YAML "key: value" lines. Group 1 = key (supports hyphens for keys like allowed-tools), - // Group 2 = quoted value, Group 3 = unquoted value. - // Accepts single or double quotes; the lazy quantifier trims trailing whitespace on unquoted values. - private static readonly Regex s_yamlKeyValueRegex = new(@"^([\w-]+)\s*:\s*(?:[""'](.+?)[""']|(.+?))\s*$", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); + // Matches top-level YAML "key: value" lines. Group 1 = key (including quotes), + // Group 2 = quoted value, Group 3 = unquoted value (possibly empty for keys such as "metadata:"). + // A value may start on a later indented line, but cannot consume another top-level field. + // Retain trailing whitespace matching so block-scalar parsing handles leading blank lines as before. + private static readonly Regex s_yamlKeyValueRegex = new(@"^([\w-]+|""[\w-]+""|'[\w-]+')[ \t]*:[ \t]*(?:\r?\n(?:[ \t]*\r?\n)*[ \t]+)?(?:[""']([^\r\n]+?)[""']|([^\r\n]*?))\s*$", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); - // Matches a "metadata:" line followed by indented sub-key/value pairs. + // Matches a metadata block, allowing quotes around the root key. // Group 1 captures the entire indented block beneath the metadata key. - private static readonly Regex s_yamlMetadataBlockRegex = new(@"^metadata\s*:\s*$\n((?:[ \t]+\S.*\n?)+)", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); + private static readonly Regex s_yamlMetadataBlockRegex = new(@"^(?:metadata|""metadata""|'metadata')\s*:\s*$\n((?:[ \t]+\S.*\n?)+)", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); // Matches indented YAML "key: value" lines within a metadata block. // Group 1 = key (supports hyphens), Group 2 = quoted value, Group 3 = unquoted value. @@ -235,30 +250,59 @@ private bool TryParseFrontmatter(string content, string skillFilePath, [NotNullW string? compatibility = null; string? allowedTools = null; + // Recognized fields use exact lowercase names, with optional quotes. Reject casing variants and + // duplicates rather than silently changing which value the skill exposes. + var seenFields = new HashSet(StringComparer.Ordinal); foreach (Match kvMatch in s_yamlKeyValueRegex.Matches(yamlContent)) { string key = kvMatch.Groups[1].Value; + key = key[0] is '"' or '\'' ? key.Substring(1, key.Length - 2) : key; + + // Unknown fields are intentionally excluded from this validation for forward compatibility. + if (!s_frontmatterFieldNames.TryGetValue(key, out string? canonicalKey)) + { + continue; + } + + if (!string.Equals(key, canonicalKey, StringComparison.Ordinal)) + { + LogIncorrectlyCasedFrontmatterField(this._logger, skillFilePath, key, canonicalKey); + return false; + } + + if (!seenFields.Add(key)) + { + LogDuplicateFrontmatterField(this._logger, skillFilePath, key); + return false; + } + + // Empty declarations participate in key validation, but leave optional scalar fields unset. + if (!kvMatch.Groups[2].Success && kvMatch.Groups[3].Length == 0) + { + continue; + } + string value = kvMatch.Groups[2].Success ? kvMatch.Groups[2].Value : ParseYamlScalarValue(yamlContent, kvMatch); - if (string.Equals(key, "name", StringComparison.OrdinalIgnoreCase)) + if (string.Equals(key, "name", StringComparison.Ordinal)) { name = value; } - else if (string.Equals(key, "description", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "description", StringComparison.Ordinal)) { description = value; } - else if (string.Equals(key, "license", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "license", StringComparison.Ordinal)) { license = value; } - else if (string.Equals(key, "compatibility", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "compatibility", StringComparison.Ordinal)) { compatibility = value; } - else if (string.Equals(key, "allowed-tools", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "allowed-tools", StringComparison.Ordinal)) { allowedTools = value; } @@ -272,7 +316,14 @@ private bool TryParseFrontmatter(string content, string skillFilePath, [NotNullW metadata = []; foreach (Match kvMatch in s_yamlIndentedKeyValueRegex.Matches(metadataMatch.Groups[1].Value)) { - metadata[kvMatch.Groups[1].Value] = kvMatch.Groups[2].Success ? kvMatch.Groups[2].Value : kvMatch.Groups[3].Value; + string key = kvMatch.Groups[1].Value; + string value = kvMatch.Groups[2].Success ? kvMatch.Groups[2].Value : kvMatch.Groups[3].Value; + + // Keep the first value and key spelling using the dictionary's case-insensitive comparison. + if (!metadata.TryAdd(key, value)) + { + LogDuplicateMetadataKey(this._logger, skillFilePath, key); + } } } @@ -722,6 +773,15 @@ private static void ValidateExtensions(IEnumerable? extensions) [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}' has an invalid '{FieldName}' value: {Reason}")] private static partial void LogInvalidFieldValue(ILogger logger, string skillFilePath, string fieldName, string reason); + [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}' contains duplicate frontmatter field '{FieldName}'")] + private static partial void LogDuplicateFrontmatterField(ILogger logger, string skillFilePath, string fieldName); + + [LoggerMessage(LogLevel.Warning, "SKILL.md at '{SkillFilePath}' contains duplicate metadata key '{Key}'; keeping the first value")] + private static partial void LogDuplicateMetadataKey(ILogger logger, string skillFilePath, string key); + + [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}' uses incorrectly cased frontmatter field '{FieldName}'; expected '{ExpectedFieldName}'")] + private static partial void LogIncorrectlyCasedFrontmatterField(ILogger logger, string skillFilePath, string fieldName, string expectedFieldName); + [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}': skill name '{SkillName}' does not match parent directory name '{DirectoryName}'")] private static partial void LogNameDirectoryMismatch(ILogger logger, string skillFilePath, string skillName, string directoryName); diff --git a/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs b/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs index 418ecd9c19c..f9754c34960 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs @@ -62,6 +62,186 @@ Content B. private const int ManyFileArchiveFileCount = 60; + [Theory] + [InlineData("description: First\ndescription: >-\n Second", "\n")] + [InlineData("description: |-\n First\nDescription: Second", "\n")] + [InlineData("description: Valid\nmetadata:\n author: First\nmetadata:\n author: Second", "\n")] + [InlineData("description: Valid\nAllowed-Tools: read", "\n")] + [InlineData("description: Valid\nallowed-tools: read\n\"allowed-tools\": other", "\n")] + [InlineData("description: Valid\nallowed-tools: read\n\"allowed-tools\": other", "\r\n")] + [InlineData("description: Valid\n'allowed-tools': other\nallowed-tools: read", "\n")] + [InlineData("description: Valid\n'allowed-tools': other\nallowed-tools: read", "\r\n")] + [InlineData("description: Valid\n\"description\": Other text", "\n")] + [InlineData("'description': Other text\ndescription: Valid", "\r\n")] + [InlineData("description: Valid\n\"Allowed-Tools\": other", "\n")] + [InlineData("description: Valid\nmetadata:\n author: First\n\"metadata\":\n author: Second", "\r\n")] + [InlineData("description: Valid\nlicense: MIT\n'license':", "\n")] + [InlineData("description: Valid\ncompatibility: Any runtime\n\"compatibility\": Other runtime", "\r\n")] + [InlineData("description: Valid\n\"name\": archived-skill", "\n")] + public async Task GetSkillsAsync_AmbiguousArchiveFrontmatter_SkipsSkillAsync(string fields, string newline) + { + // Arrange + string content = $"---\nname: archived-skill\n{fields}\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + + [Theory] + [InlineData("\n")] + [InlineData("\r\n")] + public async Task GetSkillsAsync_ArchiveValueOnIndentedNextLine_PreservedAsync(string newline) + { + // Arrange + const string Content = "---\nname: archived-skill\ndescription:\n 'Read files'\nlicense: MIT\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", Content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Equal("MIT", frontmatter.License); + } + + [Theory] + [InlineData("license", "\n")] + [InlineData("license", "\r\n")] + [InlineData("compatibility", "\n")] + [InlineData("compatibility", "\r\n")] + [InlineData("allowed-tools", "\n")] + [InlineData("allowed-tools", "\r\n")] + public async Task GetSkillsAsync_ArchiveEmptyOptionalScalar_RemainsNullAsync(string field, string newline) + { + // Arrange + string content = $"---\nname: archived-skill\ndescription: Read files\n{field}: \t\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Null(frontmatter.License); + Assert.Null(frontmatter.Compatibility); + Assert.Null(frontmatter.AllowedTools); + } + + [Theory] + [InlineData("author", "\n", "")] + [InlineData("author", "\r\n", "")] + [InlineData("Author", "\n", "")] + [InlineData("Author", "\r\n", "")] + [InlineData("author", "\n", "'")] + [InlineData("author", "\r\n", "'")] + [InlineData("Author", "\n", "'")] + [InlineData("Author", "\r\n", "'")] + [InlineData("author", "\n", "\"")] + [InlineData("author", "\r\n", "\"")] + [InlineData("Author", "\n", "\"")] + [InlineData("Author", "\r\n", "\"")] + public async Task GetSkillsAsync_ArchiveDuplicateMetadata_KeepsFirstValueAndLogsWarningAsync( + string secondKey, string newline, string quote) + { + // Arrange + string content = + $"---\n{quote}name{quote}: archived-skill\n{quote}description{quote}: >-\n Read\n files\n" + + $"{quote}compatibility{quote}:\n Any runtime\n{quote}allowed-tools{quote}: 'read'\n{quote}metadata{quote}:\n" + + $" author:\n 'First'\n {secondKey}: Second\n author: Third\n version: '1.0'\n" + + $"{quote}license{quote}: |-\n MIT\n License\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentMcpSkillsSource( + client, new() { ArchiveSkillsDirectory = this._extractionRoot }, loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var skill = Assert.Single(skills); + Assert.NotNull(skill.Frontmatter.Metadata); + Assert.Equal(2, skill.Frontmatter.Metadata.Count); + Assert.Equal("First", skill.Frontmatter.Metadata["author"]); + Assert.Equal("First", skill.Frontmatter.Metadata["Author"]); + Assert.Equal("1.0", skill.Frontmatter.Metadata["version"]); + Assert.Equal("Read files", skill.Frontmatter.Description); + Assert.Equal("MIT\nLicense", skill.Frontmatter.License); + Assert.Equal("Any runtime", skill.Frontmatter.Compatibility); + Assert.Equal("read", skill.Frontmatter.AllowedTools); + Assert.Contains("Body.", await skill.GetContentAsync()); + logger.Verify( + l => l.Log( + LogLevel.Warning, + It.IsAny(), + It.Is((state, _) => + state.ToString()!.Contains("duplicate metadata key", StringComparison.Ordinal) && + state.ToString()!.Contains("keeping the first value", StringComparison.Ordinal)), + It.IsAny(), + It.IsAny>()), + Times.Exactly(2)); + } + + [Theory] + [InlineData("First", "First", "\n")] + [InlineData("First", "First", "\r\n")] + [InlineData("\n 'First'", "First", "\n")] + [InlineData("\n 'First'", "First", "\r\n")] + [InlineData("|-\n First\n Second", "|-", "\n")] + [InlineData("|-\n First\n Second", "|-", "\r\n")] + [InlineData(">-\n First\n Second", ">-", "\n")] + [InlineData(">-\n First\n Second", ">-", "\r\n")] + public async Task GetSkillsAsync_ArchiveExistingMetadataScalarFormats_PreservedAsync(string value, string expected, string newline) + { + // Arrange + string content = "---\nname: archived-skill\ndescription: Read files\n" + + $"metadata:\n author: {value}\n version: '1.0'\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert - preserve the existing metadata scalar representation. + var skill = Assert.Single(skills); + Assert.NotNull(skill.Frontmatter.Metadata); + Assert.Equal(2, skill.Frontmatter.Metadata.Count); + Assert.Equal(expected, skill.Frontmatter.Metadata["author"]); + Assert.Equal("1.0", skill.Frontmatter.Metadata["version"]); + Assert.Contains("Body.", await skill.GetContentAsync()); + } + /// /// Malformed, unsupported, and mismatched archive digests that must be rejected. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs index 6db43554f84..1599023e1b7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs @@ -12,6 +12,8 @@ using System.Security.AccessControl; using System.Security.Principal; using System.Threading.Tasks; +using Microsoft.Extensions.Logging; +using Moq; namespace Microsoft.Agents.AI.UnitTests.AgentSkills; @@ -41,6 +43,446 @@ public void Dispose() } } + /// + /// Existing scalar representations with metadata before or after them, with and without duplicate metadata. + /// + public static TheoryData ExistingScalarFormats + { + get + { + var data = new TheoryData(); + // Pin the existing lightweight parser's output, not full YAML conformance. + (string Value, string Expected)[] cases = + [ + ("Read files", "Read files"), + ("'Read files'", "Read files"), + ("\"Read files\"", "Read files"), + ("\" Read files \"", " Read files "), + ("\"Use #tags: safely\"", "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", "Use \"quotes\" safely"), + ("\"| not a block\"", "| not a block"), + ("\n Read files", "Read files"), + ("\n\n 'Read files'", "Read files"), + ("\n >-\n Read\n files", "Read files"), + ("|\n Read\n files", "Read\nfiles"), + ("|-\n Read\n files", "Read\nfiles"), + ("|+\n Read\n files", "Read\nfiles\n"), + (">\n Read\n files", "Read files"), + (">-\n Read\n files", "Read files"), + (">+\n Read\n files", "Read files\n"), + ("|-\n\n Read\n files", "Read\nfiles"), + ("|-\n Read\n\n files", "Read\n\nfiles"), + ("|-\n Read\n indented\n files", "Read\n indented\nfiles"), + ("|-\n description: text\n Description: text", "description: text\nDescription: text"), + ("|-\n \"description\": text\n 'allowed-tools': text", "\"description\": text\n'allowed-tools': text"), + ]; + foreach (var (value, expected) in cases) + { + foreach (string newline in new[] { "\n", "\r\n" }) + { + foreach (bool metadataFirst in new[] { false, true }) + { + foreach (string quote in new[] { "", "'", "\"" }) + { + data.Add(value, expected, newline, metadataFirst, false, quote); + data.Add(value, expected, newline, metadataFirst, true, quote); + } + } + } + } + + return data; + } + } + + /// + /// Duplicate and incorrectly cased root properties with quoted names. + /// + public static TheoryData ConflictingQuotedFields + { + get + { + var data = new TheoryData(); + (string Key, string Value)[] fields = + [ + ("name", "test-skill"), + ("description", "Read files"), + ("license", "MIT"), + ("compatibility", "Any runtime"), + ("metadata", "\n author: First"), + ("allowed-tools", "read"), + ]; + foreach (var (key, value) in fields) + { + string otherFields = string.Join("\n", fields.Where(f => f.Key != key).Select(f => $"{f.Key}: {f.Value}")); + foreach (char quote in new[] { '\'', '"' }) + { + string quotedKey = $"{quote}{key}{quote}"; + string quotedField = $"{quotedKey}: {value}"; + string bareField = $"{key}: {value}"; + char otherQuote = quote == '\'' ? '"' : '\''; + foreach (string newline in new[] { "\n", "\r\n" }) + { + string incorrectlyCasedField = $"{quote}{key.ToUpperInvariant()}{quote}: {value}"; + string content = $"---\n{otherFields}\n{incorrectlyCasedField}\n---\nBody."; + string diagnostic = $"incorrectly cased frontmatter field '{key.ToUpperInvariant()}'; expected '{key}'"; + data.Add(content.Replace("\n", newline), diagnostic); + foreach (string declaration in new[] + { + $"{quotedField}\n{bareField}", + $"{bareField}\n{quotedField}", + $"{quotedField}\n{quotedField}", + $"{quotedField}\n{otherQuote}{key}{otherQuote}: {value}", + $"{quotedKey}:\n{bareField}", + $"{bareField}\n{quotedKey}:", + }) + { + content = $"---\n{otherFields}\n{declaration}\n---\nBody."; + data.Add(content.Replace("\n", newline), $"duplicate frontmatter field '{key}'"); + } + } + } + } + + return data; + } + } + + /// + /// Existing metadata scalar outputs, including the lightweight parser's incomplete YAML normalization. + /// + public static TheoryData ExistingMetadataScalarFormats + { + get + { + var data = new TheoryData(); + (string Value, string Expected)[] cases = + [ + ("First", "First"), + ("'First'", "First"), + ("\"First\"", "First"), + ("\" First \"", " First "), + ("\"Use #tags: safely\"", "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", "Use \"quotes\" safely"), + ("''", "''"), + ("\"\"", "\"\""), + ("\"\\n\"", "\\n"), + ("'It''s fine'", "It''s fine"), + ("\n First", "First"), + ("\n 'First'", "First"), + ("\n \"First\"", "First"), + ("First\n Second", "First"), + ("\"First\n Second\"", "\"First"), + ("'First\n Second'", "'First"), + ("|\n First\n Second", "|"), + ("|-\n First\n Second", "|-"), + ("|+\n First\n Second", "|+"), + (">\n First\n Second", ">"), + (">-\n First\n Second", ">-"), + (">+\n First\n Second", ">+"), + ("\n |-\n First\n Second", "|-"), + ]; + foreach (var (value, expected) in cases) + { + data.Add(value, expected, "\n"); + data.Add(value, expected, "\r\n"); + } + + return data; + } + } + + /// + /// Duplicate and incorrectly cased fields before and after inline or block values. + /// + public static TheoryData ConflictingScalarFormats + { + get + { + var data = new TheoryData(); + string[] firstValues = ["First", "'First'", "\"First\"", "|-\n First", ">-\n First"]; + string[] secondValues = ["Second", "'Second'", "\"Second\"", "|-\n Second", ">-\n Second"]; + string[] secondKeys = ["description", "Description", "DESCRIPTION"]; + foreach (string firstValue in firstValues) + { + foreach (string secondValue in secondValues) + { + foreach (string secondKey in secondKeys) + { + data.Add(firstValue, secondValue, secondKey); + } + } + } + + return data; + } + } + + /// + /// Duplicate metadata values and case-insensitive key collisions, with LF and CRLF. + /// + public static TheoryData DuplicateMetadataEntries + { + get + { + var data = new TheoryData(); + (string Key, string DuplicateKey, string Value, string DuplicateValue, string Expected)[] cases = + [ + ("author", "author", "First", "Second", "First"), + ("Author", "Author", "'First'", "\"Second\"", "First"), + ("author", "author", "\"First\"", "'Second'", "First"), + ("author", "author", "Same", "Same", "Same"), + ("author", "author", "''", "Second", "''"), + ("author", "author", "\"\"", "Second", "\"\""), + ("author", "author", "\" First \"", "Second", " First "), + ("author", "author", "0", "Second", "0"), + ("description", "description", "First", "Second", "First"), + ("metadata", "metadata", "First", "Second", "First"), + ("vendor-key", "vendor-key", "First", "Second", "First"), + ("vendor_key", "vendor_key", "First", "Second", "First"), + ("author", "Author", "First", "Second", "First"), + ("Author", "author", "First", "Second", "First"), + ("author", "AUTHOR", "First", "Second", "First"), + ("description", "Description", "First", "Second", "First"), + ("metadata", "Metadata", "First", "Second", "First"), + ("vendor-key", "Vendor-Key", "First", "Second", "First"), + ("author", "author", "\n 'First'", "Second", "First"), + ("author", "Author", "First", "\n \"Second\"", "First"), + ("author", "author", "\"author: First\"", "Second", "author: First"), + ("author", "author", "First\n # author: Not a field", "Second", "First"), + ("author", "author", "|-\n First\n paragraph", "Second", "|-"), + ("author", "Author", "First", ">-\n Second\n paragraph", "First"), + ("author", "author", ">\n First\n paragraph", "|\n Second\n paragraph", ">"), + ]; + foreach (var (key, duplicateKey, value, duplicateValue, expected) in cases) + { + data.Add(key, duplicateKey, value, duplicateValue, expected, "\n"); + data.Add(key, duplicateKey, value, duplicateValue, expected, "\r\n"); + } + + return data; + } + } + + [Theory] + [MemberData(nameof(ExistingScalarFormats))] + public async Task GetSkillsAsync_ExistingScalarFormats_PreservedAsync( + string value, string expected, string newline, bool metadataFirst, bool duplicateMetadata, string quote) + { + // Arrange + string fields = $"{quote}name{quote}: test-skill\n{quote}description{quote}: {value}\n" + + $"{quote}license{quote}: 'MIT'\n{quote}compatibility{quote}: Any runtime\n{quote}allowed-tools{quote}: read\n"; + string metadata = $"{quote}metadata{quote}:\n author: test\n" + (duplicateMetadata ? " author: Ignored\n" : ""); + string content = "---\n" + (metadataFirst ? metadata + fields : fields + metadata) + "---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal(expected, frontmatter.Description); + Assert.Equal("MIT", frontmatter.License); + Assert.Equal("Any runtime", frontmatter.Compatibility); + Assert.Equal("read", frontmatter.AllowedTools); + Assert.NotNull(frontmatter.Metadata); + Assert.Equal("test", frontmatter.Metadata["author"]); + } + + [Theory] + [MemberData(nameof(ExistingMetadataScalarFormats))] + public async Task GetSkillsAsync_ExistingMetadataScalarFormats_PreservedAsync(string value, string expected, string newline) + { + // Arrange + string content = + "---\nname: test-skill\ndescription: Read files\n" + + $"metadata:\n author: {value}\n version: '1.0'\nlicense: MIT\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert - preserve existing scalar parsing rather than expanding YAML support. + var skill = Assert.Single(skills); + Assert.Equal("Read files", skill.Frontmatter.Description); + Assert.Equal("MIT", skill.Frontmatter.License); + Assert.NotNull(skill.Frontmatter.Metadata); + Assert.Equal(2, skill.Frontmatter.Metadata.Count); + Assert.Equal(expected, skill.Frontmatter.Metadata["author"]); + Assert.Equal("1.0", skill.Frontmatter.Metadata["version"]); + Assert.Contains("Body.", await skill.GetContentAsync()); + logger.Verify( + l => l.Log( + It.Is(level => level >= LogLevel.Warning), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>()), + Times.Never); + } + + [Theory] + [MemberData(nameof(ConflictingScalarFormats))] + public async Task GetSkillsAsync_ConflictingScalarFormats_ExcludesSkillAndLogsErrorAsync( + string firstValue, + string secondValue, + string secondKey) + { + // Arrange + string content = $"---\nname: test-skill\ndescription: {firstValue}\n{secondKey}: {secondValue}\nlicense: MIT\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + string diagnostic = secondKey == "description" + ? "duplicate frontmatter field 'description'" + : $"incorrectly cased frontmatter field '{secondKey}'; expected 'description'"; + logger.Verify( + l => l.Log( + LogLevel.Error, + It.IsAny(), + It.Is((state, _) => state.ToString()!.IndexOf(diagnostic, StringComparison.Ordinal) >= 0), + It.IsAny(), + It.IsAny>()), + Times.Once); + } + + [Fact] + public async Task GetSkillsAsync_NestedMetadataAndUnknownFields_KeepsFirstMetadataValueAsync() + { + // Arrange + const string Content = + "---\nname: test-skill\ndescription: Read files\n" + + "# description: Not a field\n" + + "metadata:\n author: First\n author: Second\n Description: Nested text\n" + + " \"description\": Ignored\n 'metadata': Ignored\n" + + "vendor-option: First\nvendor-option: Second\n" + + "'vendor-option': Third\n\"VENDOR-OPTION\": Fourth\n---\n\"description\": Body text"; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", Content); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.NotNull(frontmatter.Metadata); + Assert.Equal(2, frontmatter.Metadata.Count); + Assert.Equal("First", frontmatter.Metadata["author"]); + Assert.Equal("Nested text", frontmatter.Metadata["Description"]); + } + + [Theory] + [MemberData(nameof(DuplicateMetadataEntries))] + public async Task GetSkillsAsync_DuplicateMetadata_KeepsFirstValueAndLogsWarningAsync( + string key, + string duplicateKey, + string value, + string duplicateValue, + string expected, + string newline) + { + // Arrange + string content = + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + + $" {key}: {value}\n other: Preserved\n {duplicateKey}: {duplicateValue}\n" + + $" {duplicateKey}: Third\n tail: Retained\nlicense: MIT\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var skill = Assert.Single(skills); + var frontmatter = skill.Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Equal("MIT", frontmatter.License); + Assert.Contains("Body.", await skill.GetContentAsync()); + Assert.NotNull(frontmatter.Metadata); + Assert.Equal(3, frontmatter.Metadata.Count); + Assert.Equal(expected, frontmatter.Metadata[key]); + Assert.Equal(expected, frontmatter.Metadata[duplicateKey]); + Assert.Equal("Preserved", frontmatter.Metadata["other"]); + Assert.Equal("Retained", frontmatter.Metadata["tail"]); + Assert.Equal(key, Assert.Single(frontmatter.Metadata.Keys, k => string.Equals(k, key, StringComparison.OrdinalIgnoreCase))); + string diagnostic = $"duplicate metadata key '{duplicateKey}'; keeping the first value"; + logger.Verify( + l => l.Log( + LogLevel.Warning, + It.IsAny(), + It.Is((state, _) => + state.ToString()!.IndexOf("SKILL.md", StringComparison.Ordinal) >= 0 && + state.ToString()!.IndexOf(diagnostic, StringComparison.Ordinal) >= 0), + It.IsAny(), + It.IsAny>()), + Times.Exactly(2)); + } + + [Theory] + [InlineData("''")] + [InlineData("\"\"")] + public async Task GetSkillsAsync_EmptyQuotedOptionalValues_KeepExistingRepresentationAsync(string value) + { + // Arrange + string content = + "---\nname: test-skill\ndescription: Read files\n" + + $"license: {value}\ncompatibility: {value}\nallowed-tools: {value}\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert - empty-quote normalization belongs to the separate YAML conformance work. + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal(value, frontmatter.License); + Assert.Equal(value, frontmatter.Compatibility); + Assert.Equal(value, frontmatter.AllowedTools); + } + + [Theory] + [InlineData("|")] + [InlineData("|-")] + [InlineData("|+")] + [InlineData(">")] + [InlineData(">-")] + [InlineData(">+")] + public async Task GetSkillsAsync_EmptyBlockBeforeAnotherField_ExcludesSkillAsync(string value) + { + // Arrange + _ = this.CreateSkillDirectoryWithRawContent( + "test-skill", $"---\nname: test-skill\ndescription: {value}\n\nlicense: MIT\n---\nBody."); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + [Fact] public async Task GetSkillsAsync_ValidSkill_ReturnsSkillAsync() { @@ -192,6 +634,211 @@ public async Task GetSkillsAsync_MissingDescriptionField_ExcludesSkillAsync() Assert.Empty(skills); } + [Theory] + [InlineData("name")] + [InlineData("description")] + [InlineData("license")] + [InlineData("compatibility")] + [InlineData("metadata")] + [InlineData("allowed-tools")] + public async Task GetSkillsAsync_DuplicateRecognizedFrontmatterField_ExcludesSkillAsync(string field) + { + // Arrange + string fields = field switch + { + "name" => "name: ambiguous-skill\nname: ambiguous-skill\ndescription: A skill", + "description" => "name: ambiguous-skill\ndescription: A skill\ndescription: A second description", + "metadata" => "name: ambiguous-skill\ndescription: A skill\nmetadata:\n author: first\nmetadata:\n version: 1.0", + _ => $"name: ambiguous-skill\ndescription: A skill\n{field}: first\n{field}: second", + }; + _ = this.CreateSkillDirectoryWithRawContent( + "ambiguous-skill", + $"---\n{fields}\n---\nBody."); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + + [Theory] + [InlineData("name", "Name")] + [InlineData("description", "Description")] + [InlineData("license", "License")] + [InlineData("compatibility", "Compatibility")] + [InlineData("metadata", "Metadata")] + [InlineData("allowed-tools", "Allowed-Tools")] + public async Task GetSkillsAsync_IncorrectlyCasedRecognizedFrontmatterField_ExcludesSkillAsync( + string field, + string incorrectlyCasedField) + { + // Arrange + string fields = field switch + { + "name" => $"{incorrectlyCasedField}: ambiguous-skill\ndescription: A skill", + "description" => $"name: ambiguous-skill\n{incorrectlyCasedField}: A skill", + "metadata" => $"name: ambiguous-skill\ndescription: A skill\n{incorrectlyCasedField}:\n author: test", + _ => $"name: ambiguous-skill\ndescription: A skill\n{incorrectlyCasedField}: value", + }; + _ = this.CreateSkillDirectoryWithRawContent( + "ambiguous-skill", + $"---\n{fields}\n---\nBody."); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + + [Theory] + [MemberData(nameof(ConflictingQuotedFields))] + public async Task GetSkillsAsync_ConflictingQuotedFields_ExcludesSkillAndLogsErrorAsync( + string content, + string diagnostic) + { + // Arrange + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + logger.Verify( + l => l.Log( + LogLevel.Error, + It.IsAny(), + It.Is((state, _) => state.ToString()!.IndexOf(diagnostic, StringComparison.Ordinal) >= 0), + It.IsAny(), + It.IsAny>()), + Times.Once); + } + + [Theory] + [InlineData("'", "\n")] + [InlineData("'", "\r\n")] + [InlineData("\"", "\n")] + [InlineData("\"", "\r\n")] + public async Task GetSkillsAsync_QuotedEmptyOptionalFields_RemainUnsetAsync(string quote, string newline) + { + // Arrange + string content = + $"---\n{quote}metadata{quote}:\n{quote}name{quote}: test-skill\n" + + $"{quote}description{quote}: Read files\n{quote}license{quote}:\n" + + $"{quote}compatibility{quote}:\n{quote}allowed-tools{quote}:\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("test-skill", frontmatter.Name); + Assert.Equal("Read files", frontmatter.Description); + Assert.Null(frontmatter.Metadata); + Assert.Null(frontmatter.License); + Assert.Null(frontmatter.Compatibility); + Assert.Null(frontmatter.AllowedTools); + } + + [Theory] + [InlineData("metadata", "\n")] + [InlineData("metadata", "\r\n")] + [InlineData("license", "\n")] + [InlineData("license", "\r\n")] + [InlineData("compatibility", "\n")] + [InlineData("compatibility", "\r\n")] + [InlineData("allowed-tools", "\n")] + [InlineData("allowed-tools", "\r\n")] + [InlineData("vendor-option", "\n")] + [InlineData("vendor-option", "\r\n")] + public async Task GetSkillsAsync_EmptyInlineValue_DoesNotConsumeNextFieldAsync(string field, string newline) + { + // Arrange + string content = $"---\n{field}: \t\nname: test-skill\ndescription: A test skill.\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var skill = Assert.Single(skills); + Assert.Equal("test-skill", skill.Frontmatter.Name); + Assert.Equal("A test skill.", skill.Frontmatter.Description); + Assert.Null(skill.Frontmatter.License); + Assert.Null(skill.Frontmatter.Compatibility); + Assert.Null(skill.Frontmatter.AllowedTools); + } + + [Theory] + [InlineData("license", "\n")] + [InlineData("license", "\r\n")] + [InlineData("compatibility", "\n")] + [InlineData("compatibility", "\r\n")] + [InlineData("allowed-tools", "\n")] + [InlineData("allowed-tools", "\r\n")] + public async Task GetSkillsAsync_EmptyOptionalScalarAtEnd_RemainsNullAsync(string field, string newline) + { + // Arrange + string content = $"---\nname: test-skill\ndescription: Read files\n{field}: \t\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Null(frontmatter.License); + Assert.Null(frontmatter.Compatibility); + Assert.Null(frontmatter.AllowedTools); + } + + [Theory] + [InlineData("name:", "\n")] + [InlineData("name:", "\r\n")] + [InlineData("description:", "\n")] + [InlineData("description:", "\r\n")] + [InlineData("metadata:\nmetadata:", "\n")] + [InlineData("metadata:\nmetadata:", "\r\n")] + [InlineData("license:\nlicense: MIT", "\n")] + [InlineData("license:\nlicense: MIT", "\r\n")] + [InlineData("license: MIT\nlicense:", "\n")] + [InlineData("license: MIT\nlicense:", "\r\n")] + [InlineData("compatibility:\ncompatibility: Any runtime", "\n")] + [InlineData("compatibility:\ncompatibility: Any runtime", "\r\n")] + [InlineData("allowed-tools:\nallowed-tools: read", "\n")] + [InlineData("allowed-tools:\nallowed-tools: read", "\r\n")] + [InlineData("Metadata:", "\n")] + [InlineData("Metadata:", "\r\n")] + [InlineData("allowed-tools: read\nALLOWED-TOOLS:", "\n")] + [InlineData("allowed-tools: read\nALLOWED-TOOLS:", "\r\n")] + public async Task GetSkillsAsync_EmptyInlineValue_DoesNotBypassKeyValidationAsync(string fields, string newline) + { + // Arrange + string content = $"---\n{fields}\nname: test-skill\ndescription: A test skill.\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + [Theory] [InlineData("BadName")] [InlineData("-leading-hyphen")] diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index d63912a5477..0f2bdfc5ab9 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -157,6 +157,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID ### Skills (`_skills.py`) +- **Skill frontmatter parsing** - Local and MCP archive skill loaders use PyYAML's `SafeLoader` node composition, not dictionary construction, so duplicate mapping entries remain available for validation and no YAML object constructors run. Recognized root names must be lowercase and unique after YAML decoding. Scalar fields remain text (including numeric/boolean scalar spellings); null root values remain unset. Decoded keys and retained scalar values cannot contain Unicode surrogate code points, which cannot be encoded as UTF-8. Invalid optional metadata mappings or entries warn and are skipped; case-sensitive duplicate metadata keys keep the first valid value. Invalid YAML syntax or invalid recognized root fields reject the skill. Quoting, escapes, multiline folding/chomping, and line-ending normalization follow PyYAML. YAML merge keys (`<<`) are not expanded: they are invalid at the root and skipped with warnings in metadata. - **`Skill`** - Abstract base for a skill definition bundling instructions (`content`) with frontmatter metadata, resources, and scripts. Concrete subclasses (`InlineSkill`, `FileSkill`, `ClassSkill`) accept a `frontmatter=SkillFrontmatter(...)` argument carrying the spec fields. Adding new spec fields is done in one place — on `SkillFrontmatter` — keeping the subclass constructors stable. - **`SkillFrontmatter`** - L1 discovery metadata for a skill (`name`, `description`, `license`, `compatibility`, `allowed_tools`, `metadata`). All fields are mutable plain attributes; the constructor validates `name`, `description`, and `compatibility` against the spec but post-construction assignments are not re-validated. Spec fields are reachable on every skill via `skill.frontmatter`. - **`SkillResource`** - Named supplementary content attached to a skill; holds either static `content` or a dynamic `function` (sync or async). Exactly one must be provided. diff --git a/python/packages/core/agent_framework/_skills.py b/python/packages/core/agent_framework/_skills.py index b02ec2cf5e5..eb29d6ab5b4 100644 --- a/python/packages/core/agent_framework/_skills.py +++ b/python/packages/core/agent_framework/_skills.py @@ -64,6 +64,9 @@ from pathlib import Path, PurePosixPath from typing import IO, TYPE_CHECKING, Any, ClassVar, Final, Protocol, TypeAlias, TypeVar, cast, runtime_checkable +import yaml +from yaml.nodes import MappingNode, Node, ScalarNode + from ._feature_stage import ExperimentalFeature, experimental from ._filesystem import _is_link_or_reparse_point # pyright: ignore[reportPrivateUsage] from ._sessions import ContextProvider @@ -701,7 +704,9 @@ class SkillFrontmatter: compatibility: Optional compatibility information (≤500 characters). allowed_tools: Optional space-delimited pre-approved tool names. metadata: Optional arbitrary key-value pairs (shallow-copied on - construction to avoid caller-owned dict aliasing). + construction to avoid caller-owned dict aliasing). Keys are case-sensitive. + When parsed from a SKILL.md file, invalid entries are skipped with warnings, + and exact duplicate keys retain the first valid value with warnings. """ def __init__( @@ -1749,125 +1754,29 @@ def __call__( # Matches YAML frontmatter delimited by "---" lines. # The \uFEFF? prefix allows an optional UTF-8 BOM. FRONTMATTER_RE = re.compile( - r"\A\uFEFF?---\s*$(.+?)^---\s*$", + r"\A\uFEFF?---[ \t]*\r?\n(.*?)^---[ \t]*\r?$", re.MULTILINE | re.DOTALL, ) -# Matches top-level YAML "key: value" lines (unindented). Group 1 = key, -# Group 2 = quoted value, Group 3 = unquoted value. Only matches keys at -# column 0 so that indented children (e.g. under "metadata:") are not -# mistakenly captured as top-level fields. -YAML_KV_RE = re.compile( - r"^([\w-]+)\s*:\s*(?:[\"'](.+?)[\"']|(.+?))\s*$", - re.MULTILINE, -) +FRONTMATTER_FIELD_NAMES = { + "name": "name", + "description": "description", + "license": "license", + "compatibility": "compatibility", + "metadata": "metadata", + "allowed-tools": "allowed-tools", +} -# Matches a YAML "metadata:" block followed by indented key-value pairs. -YAML_METADATA_BLOCK_RE = re.compile( - r"^metadata\s*:\s*$\n((?:[ \t]+\S.*\n?)+)", - re.MULTILINE, -) - -# Matches indented "key: value" lines within a metadata block. -YAML_INDENTED_KV_RE = re.compile( - r"^\s+([\w-]+)\s*:\s*(?:[\"'](.+?)[\"']|(.+?))\s*$", - re.MULTILINE, -) +_YAML_TEXT_SCALAR_TAGS = frozenset(f"tag:yaml.org,2002:{kind}" for kind in ("str", "bool", "int", "float", "timestamp")) # Validates skill names: lowercase letters, numbers, hyphens only; # must not start or end with a hyphen, and must not contain consecutive hyphens. VALID_NAME_RE = re.compile(r"^[a-z0-9]([a-z0-9]*-[a-z0-9])*[a-z0-9]*$") -# Block scalar indicator characters recognised by the lightweight YAML parser. -_BLOCK_SCALAR_INDICATORS = ("|", ">") - - -def _parse_yaml_scalar_value(yaml_content: str, kv_match: re.Match[str]) -> str: - """Resolve the scalar value for an unquoted YAML key-value match. - - If the captured value starts with a YAML block scalar indicator (``|`` or - ``>``), the function reads subsequent indented continuation lines, strips - the common leading indentation, and joins them according to the scalar - style (literal preserves newlines, folded replaces them with spaces). - - Chomping indicators are respected per YAML 1.2 §8.1.1.2: - - * ``-`` (strip) — final line break and trailing empty lines excluded - * ``+`` (keep) — final line break and any trailing empty lines preserved - * default (clip) — final line break preserved, trailing empty lines excluded - - For plain (non-block-scalar) values the captured text is returned as-is. - Note: explicit indentation indicators (e.g. ``|2``) are not supported; - indentation is auto-detected from the common leading whitespace. - """ - value: str = kv_match.group(3) - - if not value or value[0] not in _BLOCK_SCALAR_INDICATORS: - return value - - scalar_style = value[0] - keep_trailing_newline = len(value) > 1 and value[1] == "+" - strip_trailing_newline = len(value) > 1 and value[1] == "-" - - # Find the start of the next line after this key-value match. - next_line_start = yaml_content.find("\n", kv_match.end()) - if next_line_start < 0: - return value - next_line_start += 1 # skip the newline character itself - - # Collect indented continuation lines (or blank lines within the block). - block_lines: list[str] = [] - pos = next_line_start - while pos < len(yaml_content): - line_end = yaml_content.find("\n", pos) - if line_end < 0: - line = yaml_content[pos:] - line_end = len(yaml_content) - else: - line = yaml_content[pos:line_end] - - if not line or line.isspace(): - # Blank / whitespace-only lines are part of the block. - block_lines.append("") - pos = line_end + 1 if line_end < len(yaml_content) else line_end - continue - - if line[0] not in (" ", "\t"): - # Non-indented, non-blank line — end of the block. - break - - block_lines.append(line) - pos = line_end + 1 if line_end < len(yaml_content) else line_end - - # Strip trailing blank lines collected from the block. - while block_lines and block_lines[-1] == "": - block_lines.pop() - if not block_lines: - return "" - - # Determine the common leading indentation across non-empty lines. - # Only space/tab characters count as indentation (matches YAML semantics). - def _indent_width(s: str) -> int: - i = 0 - while i < len(s) and s[i] in (" ", "\t"): - i += 1 - return i - - common_indent = min(_indent_width(line) for line in block_lines if line) - normalized = [line[common_indent:] if line else "" for line in block_lines] - - # Literal preserves newlines; folded joins non-empty lines with spaces. - parsed = "\n".join(normalized) if scalar_style == "|" else " ".join(line for line in normalized if line) - - if keep_trailing_newline: - return parsed + "\n" - if strip_trailing_newline: - return parsed - # Clip (default): literal gets a trailing newline, folded does not. - if scalar_style == "|": - return parsed + "\n" - return parsed +def _contains_surrogate_code_points(value: str) -> bool: + """Detect surrogate escapes that PyYAML accepts but UTF-8 cannot encode.""" + return any("\ud800" <= char <= "\udfff" for char in value) # Default system prompt template for advertising available skills to the model. @@ -3600,6 +3509,58 @@ def _validate_skill_metadata( return None + @staticmethod + def _extract_frontmatter_metadata(node: Node | None, skill_file_path: str) -> dict[str, str] | None: + """Read optional metadata without rejecting a skill for invalid entries.""" + # Missing metadata, an empty declaration, or an explicit YAML null leaves metadata unset. + if node is None or (isinstance(node, ScalarNode) and node.tag == "tag:yaml.org,2002:null"): + return None + + # Metadata must be a standard YAML mapping; otherwise warn and ignore it without rejecting the skill. + if not isinstance(node, MappingNode) or node.tag != "tag:yaml.org,2002:map": + logger.warning("SKILL.md at '%s' has invalid metadata; expected a mapping", skill_file_path) + return None + + metadata: dict[str, str] = {} + + for key_node, value_node in node.value: + # Accept scalar keys as text, but reject collections, nulls, and unsupported YAML types. + if ( + not isinstance(key_node, ScalarNode) + or key_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(key_node.value) + ): + logger.warning("SKILL.md at '%s' has an invalid metadata key; skipping entry", skill_file_path) + continue + + key = key_node.value + # Keep the first valid value for each exact, case-sensitive key; warn about later duplicates. + if key in metadata: + logger.warning( + "SKILL.md at '%s' contains duplicate metadata key '%s'; keeping the first value", + skill_file_path, + key, + ) + continue + + # The Agent Skills spec requires string-to-string metadata, not arrays or nested mappings. + # Keep supported scalars as text; warn and skip invalid values rather than rejecting the skill. + if ( + not isinstance(value_node, ScalarNode) + or value_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(value_node.value) + ): + logger.warning( + "SKILL.md at '%s' has an invalid metadata value for '%s'; expected a text scalar; skipping entry", + skill_file_path, + key, + ) + continue + + metadata[key] = value_node.value + + return metadata + @staticmethod def _extract_frontmatter( content: str, @@ -3610,7 +3571,13 @@ def _extract_frontmatter( Parses the ``---``-delimited frontmatter block for all `agentskills.io specification `_ fields: ``name``, ``description``, ``license``, ``compatibility``, - ``allowed-tools``, and ``metadata``. + ``allowed-tools``, and ``metadata``. Recognized top-level fields must + use lowercase names and must not be repeated, including equivalent quoted + or escaped spellings. YAML scalars are retained as text; null root values + remain unset. Within the optional metadata mapping, keys are case-sensitive; + exact duplicates retain the first valid value and produce warnings. + Invalid metadata entries are skipped with warnings. Invalid YAML syntax + or invalid recognized root fields reject the skill. Args: content: Raw text content of the SKILL.md file. @@ -3625,41 +3592,81 @@ def _extract_frontmatter( logger.error("SKILL.md at '%s' does not contain valid YAML frontmatter delimited by '---'", skill_file_path) return None - yaml_content = match.group(1).strip() - name: str | None = None - description: str | None = None - license_value: str | None = None - compatibility: str | None = None - allowed_tools: str | None = None - - for kv_match in YAML_KV_RE.finditer(yaml_content): - key = kv_match.group(1) - value = ( - kv_match.group(2) if kv_match.group(2) is not None else _parse_yaml_scalar_value(yaml_content, kv_match) - ) + try: + # compose() keeps duplicate keys, unlike safe_load(), which silently keeps only the last value. + # The Pyright ignores are needed because PyYAML's type definitions do not fully describe compose(). + root = yaml.compose(match.group(1), Loader=yaml.SafeLoader) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + except (yaml.YAMLError, RecursionError, ValueError, OverflowError) as exc: + # PyYAML raises ValueError or OverflowError for Unicode escapes outside the valid code-point range. + logger.error("SKILL.md at '%s' contains invalid YAML frontmatter: %s", skill_file_path, exc) + return None + + if not isinstance(root, MappingNode) or root.tag != "tag:yaml.org,2002:map": + logger.error("SKILL.md at '%s' must contain a YAML frontmatter mapping", skill_file_path) + return None + + fields: dict[str, str] = {} + metadata_node: Node | None = None + seen_fields: set[str] = set() + + for key_node, value_node in root.value: + # Accept scalar keys as text, but reject collections, nulls, and unsupported YAML types. + if ( + not isinstance(key_node, ScalarNode) + or key_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(key_node.value) + ): + logger.error("SKILL.md at '%s' has an invalid frontmatter property name", skill_file_path) + return None + + key = key_node.value + canonical_key = FRONTMATTER_FIELD_NAMES.get(key.lower()) + + # Unknown fields are intentionally excluded from this validation for forward compatibility. + if canonical_key is None: + continue + + if key != canonical_key: + logger.error( + "SKILL.md at '%s' uses incorrectly cased frontmatter field '%s'; expected '%s'", + skill_file_path, + key, + canonical_key, + ) + return None + + if key in seen_fields: + logger.error("SKILL.md at '%s' contains duplicate frontmatter field '%s'", skill_file_path, key) + return None + + seen_fields.add(key) + + if key == "metadata": + metadata_node = value_node + continue + + # Leave empty declarations and explicit YAML nulls unset; required fields are validated below. + if isinstance(value_node, ScalarNode) and value_node.tag == "tag:yaml.org,2002:null": + continue + + # Reject lists, mappings, and unsupported YAML types because these fields are stored as text. + if ( + not isinstance(value_node, ScalarNode) + or value_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(value_node.value) + ): + logger.error( + "SKILL.md at '%s' has an invalid '%s' value; expected a text scalar", + skill_file_path, + key, + ) + return None - key_lower = key.lower() - if key_lower == "name": - name = value - elif key_lower == "description": - description = value - elif key_lower == "license": - license_value = value - elif key_lower == "compatibility": - compatibility = value - elif key_lower == "allowed-tools": - allowed_tools = value - - # Parse metadata block (indented key-value pairs under "metadata:"). - metadata: dict[str, str] | None = None - metadata_match = YAML_METADATA_BLOCK_RE.search(yaml_content) - if metadata_match: - metadata = {} - for kv_match in YAML_INDENTED_KV_RE.finditer(metadata_match.group(1)): - mk = kv_match.group(1) - mv = kv_match.group(2) if kv_match.group(2) is not None else kv_match.group(3) - metadata[mk] = mv + fields[key] = value_node.value + name = fields.get("name") + description = fields.get("description") + compatibility = fields.get("compatibility") error = FileSkillsSource._validate_skill_metadata(name, description, skill_file_path, compatibility) if error: logger.error(error) @@ -3670,10 +3677,10 @@ def _extract_frontmatter( return SkillFrontmatter( name=cast(str, name), description=cast(str, description), - license=license_value, + license=fields.get("license"), compatibility=compatibility, - allowed_tools=allowed_tools, - metadata=metadata, + allowed_tools=fields.get("allowed-tools"), + metadata=FileSkillsSource._extract_frontmatter_metadata(metadata_node, skill_file_path), ) @staticmethod diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index d2280bd94da..756b4a7522a 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "pydantic>=2,<3", "python-dotenv>=1,<2", "opentelemetry-api>=1.39.0,<2", + "pyyaml>=6.0,<7.0", ] [project.optional-dependencies] @@ -71,6 +72,7 @@ dev = [ "azure-ai-agentserver-core>=2.0.0b7,<3", "graphviz>=0.20,<1", "opentelemetry-exporter-otlp>=1.39.0,<2", + "types-pyyaml>=6.0,<7.0", ] [tool.uv] diff --git a/python/packages/core/tests/core/test_mcp_skills.py b/python/packages/core/tests/core/test_mcp_skills.py index 20cb8db5857..85638a036d0 100644 --- a/python/packages/core/tests/core/test_mcp_skills.py +++ b/python/packages/core/tests/core/test_mcp_skills.py @@ -757,6 +757,161 @@ def _archive_client(index_json: str, archive_url: str, archive_bytes: bytes, mim class TestMCPSkillsSourceArchive: """Tests for archive-type skill discovery via MCPSkillsSource (in-memory).""" + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + "fields", + ( + "description: First\ndescription: >-\n Second", + "description: |-\n First\nDescription: Second", + "description: Valid\nmetadata:\n author: First\nmetadata:\n author: Second", + "description: Valid\nAllowed-Tools: read", + 'description: Valid\nallowed-tools: read\n"allowed-tools": other', + "description: Valid\n'allowed-tools': other\nallowed-tools: read", + 'description: Valid\n"description": Other text', + "'description': Other text\ndescription: Valid", + 'description: Valid\n"Allowed-Tools": other', + 'description: Valid\nmetadata:\n author: First\n"metadata":\n author: Second', + "description: Valid\nlicense: MIT\n'license':", + 'description: Valid\ncompatibility: Any runtime\n"compatibility": Other runtime', + 'description: Valid\n"name": packaged-skill', + 'description: First\n"descrip\\u0074ion": Second', + 'description: First\n"\\x44escription": Second', + "description: Valid\nmetadata: [unterminated", + 'description: "\\U00110000"', + 'description: "\\UFFFFFFFF"', + 'description: "\\uD800"', + 'description: "\\uDFFF"', + ), + ) + async def test_ambiguous_archive_frontmatter_is_skipped(self, fields: str, newline: str) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = f"---\nname: packaged-skill\n{fields}\n---\nBody." + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert skills == [] + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + async def test_archive_yaml_escapes_and_invalid_metadata( + self, newline: str, caplog: pytest.LogCaptureFixture + ) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = ( + '---\nname: packaged-skill\n"\\x64escription": "Read\\nfiles"\n' + '"\\u006detadata": {author: First, "\\u0061uthor": Second, invalid: [value]}\n---\nBody.' + ) + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.description == "Read\nfiles" + assert skills[0].frontmatter.metadata == {"author": "First"} + assert len(caplog.records) == 2 + assert all(record.levelname == "WARNING" for record in caplog.records) + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + async def test_archive_value_on_indented_next_line_is_preserved(self, newline: str) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = "---\nname: packaged-skill\ndescription:\n 'Read files'\nlicense: MIT\n---\nBody." + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.description == "Read files" + assert skills[0].frontmatter.license == "MIT" + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("field", ("license", "compatibility", "allowed-tools")) + async def test_archive_empty_optional_scalar_remains_none(self, field: str, newline: str) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = f"---\nname: packaged-skill\ndescription: Read files\n{field}: \n---\nBody." + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.description == "Read files" + assert skills[0].frontmatter.license is None + assert skills[0].frontmatter.compatibility is None + assert skills[0].frontmatter.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("second_key", ("author", "Author")) + @pytest.mark.parametrize("quote", ("", "'", '"')) + async def test_archive_duplicate_metadata_keeps_first_value_and_warns( + self, newline: str, second_key: str, quote: str, caplog: pytest.LogCaptureFixture + ) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = ( + f"---\n{quote}name{quote}: packaged-skill\n{quote}description{quote}: >-\n Read\n files\n" + f"{quote}compatibility{quote}:\n Any runtime\n{quote}allowed-tools{quote}: 'read'\n" + f"{quote}metadata{quote}:\n" + f" author:\n 'First'\n {second_key}: Second\n author: Third\n version: '1.0'\n" + f"{quote}license{quote}: |-\n MIT\n License\n---\nBody." + ) + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + expected = {"author": "First", "version": "1.0"} + if second_key == "Author": + expected["Author"] = "Second" + assert skills[0].frontmatter.metadata == expected + assert skills[0].frontmatter.description == "Read files" + assert skills[0].frontmatter.license == "MIT\nLicense" + assert skills[0].frontmatter.compatibility == "Any runtime" + assert skills[0].frontmatter.allowed_tools == "read" + assert "Body." in await skills[0].get_content() + assert len(caplog.records) == (2 if second_key == "author" else 1) + assert all(record.levelname == "WARNING" for record in caplog.records) + assert all( + "duplicate metadata key 'author'; keeping the first value" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("First", "First"), + ("\n 'First'", "First"), + ("|-\n First\n Second", "First\nSecond"), + (">-\n First\n Second", "First Second"), + ), + ) + async def test_archive_yaml_metadata_scalar_formats( + self, newline: str, value: str, expected: str, caplog: pytest.LogCaptureFixture + ) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = ( + "---\nname: packaged-skill\ndescription: Read files\n" + f"metadata:\n author: {value}\n version: '1.0'\n---\nBody." + ) + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.metadata == {"author": expected, "version": "1.0"} + assert "Body." in await skills[0].get_content() + assert not caplog.records + async def test_zip_archive_discovered_as_file_skill(self) -> None: from agent_framework import FileSkill diff --git a/python/packages/core/tests/core/test_skills.py b/python/packages/core/tests/core/test_skills.py index bf6fc571b77..f15d6d93038 100644 --- a/python/packages/core/tests/core/test_skills.py +++ b/python/packages/core/tests/core/test_skills.py @@ -525,6 +525,583 @@ def test_extra_fields_parsed(self) -> None: assert result is not None assert result.name == "test-skill" + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_yaml_escaped_names_and_scalar_aliases(self, newline: str, caplog: pytest.LogCaptureFixture) -> None: + content = ( + '---\n"\\u006eame": test-skill\n"\\x64escription": &text "Read\\nfiles"\n' + 'license: *text\n"\\U0000006detadata":\n "author": First\n' + ' "\\u0061uthor": Second\n Author: Separate\n version: 1.0\n---\nBody.' + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.name == "test-skill" + assert result.description == "Read\nfiles" + assert result.license == "Read\nfiles" + assert result.metadata == {"author": "First", "Author": "Separate", "version": "1.0"} + assert len(caplog.records) == 1 + assert "duplicate metadata key 'author'" in caplog.text + + @pytest.mark.parametrize("field", ("name", "description", "license", "compatibility", "metadata", "allowed-tools")) + @pytest.mark.parametrize(("escape", "width"), (("x", 2), ("u", 4), ("U", 8))) + @pytest.mark.parametrize("uppercase", (False, True)) + def test_escaped_root_names_obey_validation( + self, field: str, escape: str, width: int, uppercase: bool, caplog: pytest.LogCaptureFixture + ) -> None: + spelling = field.upper() if uppercase else field + encoded = f'"\\{escape}{ord(spelling[0]):0{width}x}{spelling[1:]}"' + fields = { + "name": "test-skill", + "description": "Read files", + "license": "MIT", + "compatibility": "Any runtime", + "metadata": "{}", + "allowed-tools": "read", + } + content = "---\n" + "".join(f"{key}: {value}\n" for key, value in fields.items()) + content += f"{encoded}: {fields[field]}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + expected = "incorrectly cased frontmatter field" if uppercase else "duplicate frontmatter field" + assert expected in caplog.text + + def test_flow_mapping_preserves_duplicate_metadata(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\n{name: test-skill, description: Read files, " + "metadata: {author: First, author: Second, Author: Separate, enabled: true}}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First", "Author": "Separate", "enabled": "true"} + assert len(caplog.records) == 1 + assert "duplicate metadata key 'author'" in caplog.text + + def test_flow_mapping_duplicate_root_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = '---\n{name: test-skill, description: First, "\\x64escription": Second}\n---\nBody.' + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "duplicate frontmatter field 'description'" in caplog.text + + def test_aliased_root_key_duplicate_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = "---\nname: test-skill\n? &key description\n: First\n? *key\n: Second\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "duplicate frontmatter field 'description'" in caplog.text + + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("true", "true"), + ("0123", "0123"), + ("1.20", "1.20"), + ("2026-01-02", "2026-01-02"), + ("!!str null", "null"), + ('"\\u00e9"', "\u00e9"), + ('"\\U0001F600"', "\U0001f600"), + ), + ) + def test_yaml_scalar_types_remain_text(self, value: str, expected: str) -> None: + content = ( + f"---\nname: test-skill\ndescription: {value}\nlicense: {value}\ncompatibility: {value}\n" + f"allowed-tools: {value}\nmetadata:\n value: {value}\n true: boolean key\n 1: numeric key\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == expected + assert result.license == expected + assert result.compatibility == expected + assert result.allowed_tools == expected + assert result.metadata == {"value": expected, "true": "boolean key", "1": "numeric key"} + + def test_unknown_root_values_are_not_constructed(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "unknown: !!python/object/apply:builtins.str [ignored]\n" + "unknown: &recursive {self: *recursive}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert not caplog.records + + def test_root_merge_key_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = "---\nname: test-skill\ndescription: Read files\n<<: {description: Override}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid frontmatter property name" in caplog.text + + def test_metadata_merge_key_is_skipped(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "metadata: {<<: {author: Ignored}, author: First}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First"} + assert len(caplog.records) == 1 + assert "invalid metadata key; skipping entry" in caplog.text + + @pytest.mark.parametrize("value", ("", "null", "~")) + def test_yaml_null_optional_fields_remain_unset(self, value: str, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + f"metadata: {value}\nlicense: {value}\ncompatibility: {value}\nallowed-tools: {value}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata is None + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + assert not caplog.records + + @pytest.mark.parametrize("value", ("[First, Second]", "text", "42", "!!null {key: value}")) + def test_invalid_metadata_mapping_only_warns(self, value: str, caplog: pytest.LogCaptureFixture) -> None: + content = f"---\nname: test-skill\ndescription: Read files\nmetadata: {value}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata is None + assert len(caplog.records) == 1 + assert caplog.records[0].levelname == "WARNING" + assert "invalid metadata; expected a mapping" in caplog.text + + def test_invalid_metadata_entries_are_skipped(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + " valid: Keep\n mapping: {key: value}\n sequence: [First, Second]\n missing:\n" + " tagged: !!python/name:builtins.str text\n ? [complex, key]\n : value\n" + " author: {invalid: value}\n author: First valid\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"valid": "Keep", "author": "First valid"} + assert len(caplog.records) == 6 + assert all(record.levelname == "WARNING" for record in caplog.records) + assert all("skipping entry" in record.getMessage() for record in caplog.records) + + def test_recursive_metadata_alias_is_skipped(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "metadata: &metadata {recursive: *metadata, author: First}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First"} + assert len(caplog.records) == 1 + assert "invalid metadata value for 'recursive'" in caplog.text + + @pytest.mark.parametrize("escape", ("\\uD800", "\\uDFFF", "\\uD83D\\uDE00")) + @pytest.mark.parametrize("invalid_key", (False, True)) + def test_metadata_surrogates_are_skipped( + self, escape: str, invalid_key: bool, caplog: pytest.LogCaptureFixture + ) -> None: + entry = f'"{escape}": Ignored' if invalid_key else f'author: "{escape}"' + content = f"---\nname: test-skill\ndescription: Read files\nmetadata:\n {entry}\n author: First valid\n---" + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First valid"} + assert len(caplog.records) == 1 + assert caplog.records[0].levelname == "WARNING" + assert ("invalid metadata key" if invalid_key else "invalid metadata value") in caplog.text + + @pytest.mark.parametrize("field", ("name", "description", "license", "compatibility", "allowed-tools")) + @pytest.mark.parametrize( + "value", + ( + "[First, Second]", + "{key: value}", + "!!null [value]", + "!!python/name:builtins.str text", + '"\\uD800"', + '"\\uDFFF"', + '"\\uD83D\\uDE00"', + ), + ) + def test_non_text_root_values_rejected(self, field: str, value: str, caplog: pytest.LogCaptureFixture) -> None: + fields = {"name": "test-skill", "description": "Read files", field: value} + content = "---\n" + "".join(f"{key}: {item}\n" for key, item in fields.items()) + "---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert f"invalid '{field}' value; expected a text scalar" in caplog.text + + @pytest.mark.parametrize("header", ("", "null", "[]", "plain text", "!!python/object/apply:builtins.str [text]")) + def test_non_mapping_frontmatter_rejected(self, header: str, caplog: pytest.LogCaptureFixture) -> None: + result = FileSkillsSource._extract_frontmatter(f"---\n{header}\n---\nBody.", "test.md") + + assert result is None + assert "must contain a YAML frontmatter mapping" in caplog.text + + @pytest.mark.parametrize( + "fields", + ( + "metadata: [unterminated", + 'license: "\\q"', + 'license: "\\U00110000"', + 'license: "\\UFFFFFFFF"', + "license: *missing", + "description: \tvalue", + ), + ) + def test_invalid_yaml_rejected(self, fields: str, caplog: pytest.LogCaptureFixture) -> None: + content = f"---\nname: test-skill\ndescription: Read files\n{fields}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid YAML frontmatter" in caplog.text + + @pytest.mark.parametrize("key", ("[complex, key]", '"\\uD800"', '"\\uDFFF"')) + def test_invalid_root_property_name_rejected(self, key: str, caplog: pytest.LogCaptureFixture) -> None: + content = f"---\nname: test-skill\ndescription: Read files\n? {key}\n: value\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid frontmatter property name" in caplog.text + + def test_excessive_yaml_nesting_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = "---\nname: test-skill\ndescription: Read files\nunknown: " + "[" * 2000 + "]" * 2000 + "\n---" + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid YAML frontmatter" in caplog.text + + @pytest.mark.parametrize( + "field", + ("name", "description", "license", "compatibility", "metadata", "allowed-tools"), + ) + def test_duplicate_recognized_field_rejected(self, field: str) -> None: + if field == "name": + fields = "name: test-skill\nname: test-skill\ndescription: A test skill." + elif field == "description": + fields = "name: test-skill\ndescription: A test skill.\ndescription: A second description." + elif field == "metadata": + fields = ( + "name: test-skill\ndescription: A test skill.\nmetadata:\n author: first\nmetadata:\n version: 1.0" + ) + else: + fields = f"name: test-skill\ndescription: A test skill.\n{field}: first\n{field}: second" + + result = FileSkillsSource._extract_frontmatter(f"---\n{fields}\n---\nBody.", "test.md") + + assert result is None + + @pytest.mark.parametrize( + ("field", "incorrect_casing"), + ( + ("name", "Name"), + ("description", "Description"), + ("license", "License"), + ("compatibility", "Compatibility"), + ("metadata", "Metadata"), + ("allowed-tools", "Allowed-Tools"), + ), + ) + def test_incorrectly_cased_recognized_field_rejected(self, field: str, incorrect_casing: str) -> None: + if field == "name": + fields = f"{incorrect_casing}: test-skill\ndescription: A test skill." + elif field == "description": + fields = f"name: test-skill\n{incorrect_casing}: A test skill." + elif field == "metadata": + fields = f"name: test-skill\ndescription: A test skill.\n{incorrect_casing}:\n author: test" + else: + fields = f"name: test-skill\ndescription: A test skill.\n{incorrect_casing}: value" + + result = FileSkillsSource._extract_frontmatter(f"---\n{fields}\n---\nBody.", "test.md") + + assert result is None + + @pytest.mark.parametrize("field", ("name", "description", "license", "compatibility", "metadata", "allowed-tools")) + @pytest.mark.parametrize("quote", ("'", '"')) + @pytest.mark.parametrize("layout", ("uppercase", "before", "after", "both", "mixed", "empty-first", "empty-last")) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_conflicting_quoted_fields_rejected( + self, + field: str, + quote: str, + layout: str, + newline: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + values = { + "name": "test-skill", + "description": "Read files", + "license": "MIT", + "compatibility": "Any runtime", + "metadata": "\n author: First", + "allowed-tools": "read", + } + fields = "\n".join(f"{key}: {value}" for key, value in values.items() if key != field) + quoted_key = f"{quote}{field}{quote}" + quoted_field = f"{quoted_key}: {values[field]}" + bare_field = f"{field}: {values[field]}" + other_quote = '"' if quote == "'" else "'" + declarations = { + "uppercase": f"{quote}{field.upper()}{quote}: {values[field]}", + "before": f"{quoted_field}\n{bare_field}", + "after": f"{bare_field}\n{quoted_field}", + "both": f"{quoted_field}\n{quoted_field}", + "mixed": f"{quoted_field}\n{other_quote}{field}{other_quote}: {values[field]}", + "empty-first": f"{quoted_key}:\n{bare_field}", + "empty-last": f"{bare_field}\n{quoted_key}:", + } + content = f"---\n{fields}\n{declarations[layout]}\n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + diagnostic = ( + f"incorrectly cased frontmatter field '{field.upper()}'; expected '{field}'" + if layout == "uppercase" + else f"duplicate frontmatter field '{field}'" + ) + assert diagnostic in caplog.text + + @pytest.mark.parametrize("quote", ("'", '"')) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_quoted_empty_optional_fields_remain_unset(self, quote: str, newline: str) -> None: + content = ( + f"---\n{quote}metadata{quote}:\n{quote}name{quote}: test-skill\n" + f"{quote}description{quote}: Read files\n{quote}license{quote}:\n" + f"{quote}compatibility{quote}:\n{quote}allowed-tools{quote}:\n---\nBody." + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.name == "test-skill" + assert result.description == "Read files" + assert result.metadata is None + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("field", ("metadata", "license", "compatibility", "allowed-tools", "vendor-option")) + def test_empty_inline_value_does_not_consume_next_field(self, field: str, newline: str) -> None: + content = newline.join(( + "---", + f"{field}: ", + "name: test-skill", + "description: A test skill.", + "---", + "Body.", + )) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.name == "test-skill" + assert result.description == "A test skill." + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("field", ("license", "compatibility", "allowed-tools")) + def test_empty_optional_scalar_at_end_remains_none(self, field: str, newline: str) -> None: + content = f"---\nname: test-skill\ndescription: Read files\n{field}: \n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + "fields", + ( + "name:", + "description:", + "metadata:\nmetadata:", + "license:\nlicense: MIT", + "license: MIT\nlicense:", + "compatibility:\ncompatibility: Any runtime", + "allowed-tools:\nallowed-tools: read", + "Metadata:", + "allowed-tools: read\nALLOWED-TOOLS:", + ), + ) + def test_empty_inline_value_does_not_bypass_key_validation(self, fields: str, newline: str) -> None: + content = f"---\n{fields}\nname: test-skill\ndescription: A test skill.\n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + + @pytest.mark.parametrize("first_value", ("First", "'First'", '"First"', "|-\n First", ">-\n First")) + @pytest.mark.parametrize("second_value", ("Second", "'Second'", '"Second"', "|-\n Second", ">-\n Second")) + @pytest.mark.parametrize("second_key", ("description", "Description", "DESCRIPTION")) + def test_duplicate_and_cased_fields_across_scalar_formats( + self, first_value: str, second_value: str, second_key: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + f"---\nname: test-skill\ndescription: {first_value}\n{second_key}: {second_value}\nlicense: MIT\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + diagnostic = ( + "duplicate frontmatter field 'description'" + if second_key == "description" + else f"incorrectly cased frontmatter field '{second_key}'; expected 'description'" + ) + assert diagnostic in caplog.text + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("key", "first_value", "second_value", "expected"), + ( + ("author", "First", "Second", "First"), + ("Author", "'First'", '"Second"', "First"), + ("author", '"First"', "'Second'", "First"), + ("author", "Same", "Same", "Same"), + ("author", "''", "Second", ""), + ("author", '""', "Second", ""), + ("author", '" First "', "Second", " First "), + ("author", "0", "Second", "0"), + ("description", "First", "Second", "First"), + ("metadata", "First", "Second", "First"), + ("vendor-key", "First", "Second", "First"), + ("vendor_key", "First", "Second", "First"), + ("author", "\n 'First'", "Second", "First"), + ("author", "First", '\n "Second"', "First"), + ("author", '"author: First"', "Second", "author: First"), + ("author", "First\n # author: Not a field", "Second", "First"), + ("author", "|-\n First\n paragraph", "Second", "First\nparagraph"), + ("author", "First", ">-\n Second\n paragraph", "First"), + ("author", ">\n First\n paragraph", "|\n Second\n paragraph", "First paragraph\n"), + ), + ) + def test_duplicate_metadata_keeps_first_value_and_warns( + self, + newline: str, + key: str, + first_value: str, + second_value: str, + expected: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + f" {key}: {first_value}\n other: Preserved\n {key}: {second_value}\n" + f" {key}: Third\n tail: Retained\nlicense: MIT\n---\nBody." + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.license == "MIT" + assert result.metadata == {key: expected, "other": "Preserved", "tail": "Retained"} + assert [(record.levelname, record.getMessage()) for record in caplog.records] == [ + ("WARNING", f"SKILL.md at 'test.md' contains duplicate metadata key '{key}'; keeping the first value") + ] * 2 + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("first_key", "second_key"), + ( + ("author", "Author"), + ("Author", "author"), + ("author", "AUTHOR"), + ("description", "Description"), + ("metadata", "Metadata"), + ("vendor-key", "Vendor-Key"), + ), + ) + def test_metadata_keys_are_case_sensitive( + self, newline: str, first_key: str, second_key: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + f" {first_key}: First\n {second_key}: Second\n---\nBody." + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.metadata == {first_key: "First", second_key: "Second"} + assert not caplog.records + + def test_nested_metadata_keeps_first_value_and_unknown_fields_are_ignored(self) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "# description: Not a field\n" + "metadata:\n author: First\n author: Second\n Description: Nested text\n" + " \"description\": Nested quoted text\n 'metadata': Nested metadata value\n" + "vendor-option: First\nvendor-option: Second\n" + '\'vendor-option\': Third\n"VENDOR-OPTION": Fourth\n---\n"description": Body text' + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.metadata == { + "author": "First", + "Description": "Nested text", + "description": "Nested quoted text", + "metadata": "Nested metadata value", + } + + @pytest.mark.parametrize("value", ("''", '""')) + def test_empty_quoted_optional_values_are_empty_strings(self, value: str) -> None: + content = ( + f"---\nname: test-skill\ndescription: Read files\n" + f"license: {value}\ncompatibility: {value}\nallowed-tools: {value}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.license == "" + assert result.compatibility == "" + assert result.allowed_tools == "" + + @pytest.mark.parametrize("value", ("|", "|-", "|+", ">", ">-", ">+")) + def test_empty_block_before_another_field_rejected(self, value: str) -> None: + content = f"---\nname: test-skill\ndescription: {value}\n\nlicense: MIT\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + # --------------------------------------------------------------------------- # Tests: skill discovery and loading @@ -534,12 +1111,212 @@ def test_extra_fields_parsed(self) -> None: class TestDiscoverAndLoadSkills: """Tests for file skill discovery via FileSkillsSource.get_skills().""" + @pytest.mark.parametrize( + ("fields", "loads"), + ( + ( + ( + '"\\x64escription": "Read\\nfiles"\n' + 'metadata: {author: First, "\\u0061uthor": Ignored, invalid: [item]}' + ), + True, + ), + ('description: Read files\n"\\x64escription": Second', False), + ('"\\x44escription": Read files', False), + ("description: Read files\nmetadata: [unterminated", False), + ('description: "\\U00110000"', False), + ('description: "\\UFFFFFFFF"', False), + ('description: "\\uD800"', False), + ('description: "\\uDFFF"', False), + ), + ) + async def test_file_yaml_decoding_and_validation( + self, tmp_path: Path, fields: str, loads: bool, caplog: pytest.LogCaptureFixture + ) -> None: + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_text(f"---\nname: test-skill\n{fields}\n---\nBody.", encoding="utf-8") + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + if loads: + assert len(skills) == 1 + assert skills["test-skill"].frontmatter.description == "Read\nfiles" + assert skills["test-skill"].frontmatter.metadata == {"author": "First"} + assert len(caplog.records) == 2 + assert all(record.levelname == "WARNING" for record in caplog.records) + else: + assert skills == {} + assert any(record.levelname == "ERROR" for record in caplog.records) + + @pytest.mark.parametrize("metadata_first", (False, True)) + @pytest.mark.parametrize("duplicate_metadata", (False, True)) + @pytest.mark.parametrize("quote", ("", "'", '"')) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("Read files", "Read files"), + ("'Read files'", "Read files"), + ('"Read files"', "Read files"), + ('" Read files "', " Read files "), + ('"Use #tags: safely"', "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", 'Use "quotes" safely'), + ('"| not a block"', "| not a block"), + ("\n Read files", "Read files"), + ("\n\n 'Read files'", "Read files"), + ("\n >-\n Read\n files", "Read files"), + ("|\n Read\n files", "Read\nfiles\n"), + ("|-\n Read\n files", "Read\nfiles"), + ("|+\n Read\n files", "Read\nfiles\n"), + (">\n Read\n files", "Read files\n"), + (">-\n Read\n files", "Read files"), + (">+\n Read\n files", "Read files\n"), + ("|-\n\n Read\n files", "\nRead\nfiles"), + ("|-\n Read\n\n files", "Read\n\nfiles"), + ("|-\n Read\n indented\n files", "Read\n indented\nfiles"), + ("|-\n description: text\n Description: text", "description: text\nDescription: text"), + ("|-\n \"description\": text\n 'allowed-tools': text", "\"description\": text\n'allowed-tools': text"), + ("Read files # comment", "Read files"), + ('"Read\\nfiles"', "Read\nfiles"), + ("Read\n files", "Read files"), + ("|2-\n Read\n files", "Read\nfiles"), + (">2-\n Read\n files", "Read files"), + ), + ) + async def test_yaml_scalar_formats( + self, + tmp_path: Path, + newline: str, + value: str, + expected: str, + metadata_first: bool, + duplicate_metadata: bool, + quote: str, + ) -> None: + fields = ( + f"{quote}name{quote}: test-skill\n{quote}description{quote}: {value}\n" + f"{quote}license{quote}: 'MIT'\n{quote}compatibility{quote}: Any runtime\n" + f"{quote}allowed-tools{quote}: read\n" + ) + metadata = f"{quote}metadata{quote}:\n author: test\n" + (" author: Ignored\n" if duplicate_metadata else "") + content = "---\n" + (metadata + fields if metadata_first else fields + metadata) + "---\nBody." + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert len(skills) == 1 + frontmatter = skills["test-skill"].frontmatter + assert frontmatter.description == expected + assert frontmatter.license == "MIT" + assert frontmatter.compatibility == "Any runtime" + assert frontmatter.allowed_tools == "read" + assert frontmatter.metadata == {"author": "test"} + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("First", "First"), + ("'First'", "First"), + ('"First"', "First"), + ('" First "', " First "), + ('"Use #tags: safely"', "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", 'Use "quotes" safely'), + ("''", ""), + ('""', ""), + ('"\\n"', "\n"), + ("'It''s fine'", "It's fine"), + ("\n First", "First"), + ("\n 'First'", "First"), + ('\n "First"', "First"), + ("First\n Second", "First Second"), + ('"First\n Second"', "First Second"), + ("'First\n Second'", "First Second"), + ("|\n First\n Second", "First\nSecond\n"), + ("|-\n First\n Second", "First\nSecond"), + ("|+\n First\n Second", "First\nSecond\n"), + (">\n First\n Second", "First Second\n"), + (">-\n First\n Second", "First Second"), + (">+\n First\n Second", "First Second\n"), + ("\n |-\n First\n Second", "First\nSecond"), + ), + ) + async def test_yaml_metadata_scalar_formats( + self, tmp_path: Path, newline: str, value: str, expected: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + f"metadata:\n author: {value}\n version: '1.0'\nlicense: MIT\n---\nBody." + ) + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert len(skills) == 1 + skill = skills["test-skill"] + assert skill.frontmatter.description == "Read files" + assert skill.frontmatter.license == "MIT" + assert skill.frontmatter.metadata == {"author": expected, "version": "1.0"} + assert "Body." in await skill.get_content() + assert not caplog.records + async def test_discovers_valid_skill(self, tmp_path: Path) -> None: _write_skill(tmp_path, "my-skill") skills = await _discover_file_skills_for_test([str(tmp_path)]) assert "my-skill" in skills assert skills["my-skill"].frontmatter.name == "my-skill" + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + "fields", + ( + 'allowed-tools: read\n"allowed-tools": other', + "'allowed-tools': other\nallowed-tools: read", + '"Description": Other text', + "metadata:\n author: First\n'metadata':\n author: Second", + ), + ) + async def test_conflicting_quoted_root_field_prevents_discovery( + self, tmp_path: Path, fields: str, newline: str + ) -> None: + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + content = f"---\nname: test-skill\ndescription: Read files\n{fields}\n---\nBody." + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert not skills + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + async def test_duplicate_metadata_does_not_prevent_discovery( + self, tmp_path: Path, newline: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + " author: First\n Author: Separate\n author: Ignored\n version: '1.0'\n---\nBody." + ) + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert len(skills) == 1 + skill = skills["test-skill"] + assert skill.frontmatter.metadata == {"author": "First", "Author": "Separate", "version": "1.0"} + assert "Body." in await skill.get_content() + assert len(caplog.records) == 1 + assert caplog.records[0].levelname == "WARNING" + assert "duplicate metadata key 'author'; keeping the first value" in caplog.text + async def test_discovers_nested_skills(self, tmp_path: Path) -> None: skills_dir = tmp_path / "skills" _write_skill(skills_dir, "skill-a") @@ -2515,11 +3292,21 @@ def test_literal_block_scalar(self) -> None: assert result is not None assert result.description == "Line one\nLine two\n" + @pytest.mark.parametrize(("indicator", "expected"), (("|-", "Read"), ("|", "Read\n"), ("|+", "Read\n\n\n"))) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_terminal_block_chomping_preserves_yaml_content(self, indicator: str, expected: str, newline: str) -> None: + content = f"---\nname: test-skill\ndescription: {indicator}\n Read\n\n\n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == expected + def test_folded_block_scalar(self) -> None: content = "---\nname: test-skill\ndescription: >\n This is a multi-line\n description block\n---\nBody." result = FileSkillsSource._extract_frontmatter(content, "test.md") assert result is not None - assert result.description == "This is a multi-line description block" + assert result.description == "This is a multi-line description block\n" def test_literal_strip_chomping(self) -> None: content = "---\nname: test-skill\ndescription: |-\n No trailing newline\n---\nBody." @@ -2583,14 +3370,14 @@ def test_folded_block_scalar_real_skill_format(self) -> None: assert result.description == ( "Coding standards, conventions, and patterns for developing Python code in the " "Agent Framework repository. Use this when writing or modifying Python source " - "files in the python/ directory." + "files in the python/ directory.\n" ) def test_block_scalar_with_other_fields_after(self) -> None: content = "---\nname: test-skill\ndescription: >\n A folded\n description\nlicense: MIT\n---\nBody." result = FileSkillsSource._extract_frontmatter(content, "test.md") assert result is not None - assert result.description == "A folded description" + assert result.description == "A folded description\n" assert result.license == "MIT" def test_plain_value_unchanged(self) -> None: @@ -2623,14 +3410,14 @@ def test_block_scalar_on_license_field(self) -> None: ) result = FileSkillsSource._extract_frontmatter(content, "test.md") assert result is not None - assert result.license == "Custom license spanning multiple lines" + assert result.license == "Custom license spanning multiple lines\n" - def test_block_scalar_tab_indentation(self) -> None: - """Tab characters should count as indentation for block scalar continuation lines.""" + def test_block_scalar_tab_indentation_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + """YAML indentation must use spaces.""" content = "---\nname: test-skill\ndescription: |\n\tTab-indented line one\n\tTab-indented line two\n---\nBody." result = FileSkillsSource._extract_frontmatter(content, "test.md") - assert result is not None - assert result.description == "Tab-indented line one\nTab-indented line two\n" + assert result is None + assert "invalid YAML frontmatter" in caplog.text def test_block_scalar_blank_line_within_block(self) -> None: """Blank lines within a block scalar should be preserved as paragraph separators.""" diff --git a/python/uv.lock b/python/uv.lock index 27bf49e5eb1..2ae5de419bd 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -449,6 +449,7 @@ dependencies = [ { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "python-dotenv", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] @@ -491,6 +492,7 @@ dev = [ { name = "azure-ai-agentserver-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "graphviz", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "opentelemetry-exporter-otlp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "types-pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] [package.metadata] @@ -529,6 +531,7 @@ requires-dist = [ { name = "opentelemetry-api", specifier = ">=1.39.0,<2" }, { name = "pydantic", specifier = ">=2,<3" }, { name = "python-dotenv", specifier = ">=1,<2" }, + { name = "pyyaml", specifier = ">=6.0,<7.0" }, { name = "typing-extensions", specifier = ">=4.15.0,<5" }, ] provides-extras = ["all"] @@ -538,6 +541,7 @@ dev = [ { name = "azure-ai-agentserver-core", specifier = ">=2.0.0b7,<3" }, { name = "graphviz", specifier = ">=0.20,<1" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.39.0,<2" }, + { name = "types-pyyaml", specifier = ">=6.0,<7.0" }, ] [[package]] From a3fc29dbd63a5223ece479b9b7eed53f50144999 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 17 Sep 2026 08:47:44 +0000 Subject: [PATCH 05/26] .NET/Python: Add more content types to purview handling (#8370) * Add more content types to purview handling * Address PR comments * Split Purview content processing requests Create one processContent request per content entry in both .NET and Python while preserving the Graph contentEntries array contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Skip empty Purview content Avoid sending empty text or binary data to Graph processContent APIs in both .NET and Python. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address PR comments --------- Co-authored-by: Eoin Doherty Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../Models/Common/ContentToProcess.cs | 6 +- .../src/Microsoft.Agents.AI.Purview/README.md | 36 ++ .../ScopedContentProcessor.cs | 229 ++++++++- .../PurviewClientTests.cs | 2 +- .../ScopedContentProcessorTests.cs | 448 ++++++++++++++++- python/packages/purview/README.md | 40 +- .../agent_framework_purview/_models.py | 16 +- .../agent_framework_purview/_processor.py | 222 ++++++-- .../purview/tests/purview/conftest.py | 2 +- .../purview/tests/purview/test_processor.py | 475 +++++++++++++++++- .../tests/purview/test_purview_models.py | 7 +- 11 files changed, 1396 insertions(+), 87 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs index 9e2e5824f3a..38bdd10b681 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs @@ -14,19 +14,19 @@ internal sealed class ContentToProcess /// /// Creates a new instance of ContentToProcess. /// - /// The content to send and its associated ids. + /// The content to send and its associated id. /// Metadata about the activity performed with the content. /// Metadata about the device that produced the content. /// Metadata about the application integrating with Purview. /// Metadata about the application being protected by Purview. public ContentToProcess( - List contentEntries, + ProcessContentMetadataBase contentEntry, ActivityMetadata activityMetadata, DeviceMetadata deviceMetadata, IntegratedAppMetadata integratedAppMetadata, ProtectedAppMetadata protectedAppMetadata) { - this.ContentEntries = contentEntries; + this.ContentEntries = [contentEntry]; this.ActivityMetadata = activityMetadata; this.DeviceMetadata = deviceMetadata; this.IntegratedAppMetadata = integratedAppMetadata; diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/README.md b/dotnet/src/Microsoft.Agents.AI.Purview/README.md index 39915507b25..e92d1380e13 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/README.md +++ b/dotnet/src/Microsoft.Agents.AI.Purview/README.md @@ -273,3 +273,39 @@ catch (PurviewException e) this._logger.LogError(e, "Purview middleware threw an exception.") } ``` + +## Security Considerations + +### Identity is a trusted input + +Purview evaluates DLP policy **for a specific user**. The identity this integration resolves therefore +decides *which* policy is applied, and it is resolved in this order: + +1. The user id from the configured `TokenCredential`'s token, when the credential resolves to a user. +2. The `userId` argument passed to the processor. +3. `ChatMessage.AdditionalProperties["userId"]`. +4. `ChatMessage.AuthorName`, when it is a GUID. + +Only source 1 is verified. Sources 2-4 are supplied by the hosting application, so **a host must not +populate them from data that has crossed a trust boundary**. If an end user, an upstream service or a +model response can influence `AdditionalProperties["userId"]` or `AuthorName`, that party can select a +different user's DLP policy - typically one with weaker rules - and evade enforcement. Where identity +must come from a request, derive it from a validated token on the server, never from the request body. +Prefer a user-delegated credential (source 1) whenever possible. + +`PurviewAppLocation` in `PurviewSettings` is trusted in the same way: it selects which policy locations +apply and must be configured by the host, not by the caller. + +### Fail-closed behaviour + +Policy evaluation fails closed. If no user id can be resolved, or the tenant or app location cannot be +determined, `PurviewRequestException` is thrown rather than letting content through unevaluated. Use +`IgnoreExceptions` if you deliberately want availability over enforcement - but understand that it +disables enforcement for every error, not just transient ones. + +### What is evaluated + +Every content item on a message is submitted for evaluation, not just its text: `DataContent` is sent +as Purview binary content, and `FunctionCallContent`, `FunctionResultContent` and other structured +content are serialized to text. Only `UsageContent` is skipped, because it carries token counts rather +than user data. diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs index f49f1239a89..a456f2d7891 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; @@ -22,6 +23,11 @@ internal sealed class ScopedContentProcessor : IScopedContentProcessor private readonly ICacheProvider _cacheProvider; private readonly IChannelHandler _channelHandler; + /// + /// The characters that terminate the authority of a URL, being the start of the path, query or fragment. + /// + private static readonly char[] s_authorityDelimiters = ['/', '?', '#']; + /// /// Create a new instance of . /// @@ -52,12 +58,7 @@ public ScopedContentProcessor(IPurviewClient purviewClient, ICacheProvider cache foreach (DlpActionInfo policyAction in processContentResponse.PolicyActions) { // We need to process all data before blocking, so set the flag and return it outside of this loop. - if (policyAction.Action == DlpAction.BlockAccess) - { - shouldBlock = true; - } - - if (policyAction.RestrictionAction == RestrictionAction.Block) + if (IsBlockingAction(policyAction)) { shouldBlock = true; } @@ -68,6 +69,85 @@ public ScopedContentProcessor(IPurviewClient purviewClient, ICacheProvider cache return (shouldBlock, resolvedUserId); } + /// + /// Whether a policy action means the content must not be released. + /// + /// + /// is not blocking on its own: it carries a separate + /// that selects the enforcement mode, which may be + /// , or + /// as well as . Only an + /// explicit mode withholds the content. + /// + /// The policy action to inspect. + /// when the action blocks the content. + private static bool IsBlockingAction(DlpActionInfo actionInfo) + => actionInfo.Action == DlpAction.BlockAccess + || actionInfo.RestrictionAction == RestrictionAction.Block; + + /// + /// Filter policy actions down to the ones that block. + /// + /// The policy actions to filter. + /// The blocking subset of . + private static List GetBlockingActions(List actionInfos) + { + List blockingActions = []; + + foreach (DlpActionInfo actionInfo in actionInfos) + { + if (IsBlockingAction(actionInfo)) + { + blockingActions.Add(actionInfo); + } + } + + return blockingActions; + } + + /// + /// Normalize a policy location value for comparison, according to its location type. + /// + /// + /// Application ids (GUIDs) and domain names are case-insensitive, so those fold whole. URL values + /// are not. Only the scheme and the host are case-insensitive; the userinfo, path, query and + /// fragment are all case-sensitive, so folding a URL whole would let a scope for + /// contoso.com/public match a request for contoso.com/Public. Location types that are + /// not recognized fold whole, which matches more scopes rather than fewer. + /// + /// The location type segment, for example policyLocationUrl. + /// The location value to normalize. + /// The value in its comparable form. + private static string NormalizeLocationValue(string locationType, string value) + { + if (!locationType.EndsWith("url", StringComparison.OrdinalIgnoreCase)) + { + return value.ToUpperInvariant(); + } + + int schemeEnd = value.IndexOf("://", StringComparison.Ordinal); + int authorityStart = schemeEnd >= 0 ? schemeEnd + 3 : 0; + + // The authority ends at the first path, query or fragment delimiter, whichever comes first. + int authorityEnd = value.IndexOfAny(s_authorityDelimiters, authorityStart); + if (authorityEnd < 0) + { + authorityEnd = value.Length; + } + + // Credentials are case-sensitive, so only the host half of the authority folds. + string authority = value.Substring(authorityStart, authorityEnd - authorityStart); + int userInfoEnd = authority.LastIndexOf('@'); + string userInfo = userInfoEnd >= 0 ? authority.Substring(0, userInfoEnd + 1) : string.Empty; + string host = userInfoEnd >= 0 ? authority.Substring(userInfoEnd + 1) : authority; + + return string.Concat( + value.Substring(0, authorityStart).ToUpperInvariant(), + userInfo, + host.ToUpperInvariant(), + value.Substring(authorityEnd)); + } + private static bool TryGetUserIdFromPayload(IEnumerable messages, out string? userId) { userId = null; @@ -115,12 +195,9 @@ private async Task> MapMessageToPCRequestsAsync(IEnu foreach (ChatMessage message in messages) { string messageId = message.MessageId ?? Guid.NewGuid().ToString(); - ContentBase content = new PurviewTextContent(message.Text); string correlationId = (sessionId ?? Guid.NewGuid().ToString()) + "@AF"; - ProcessConversationMetadata conversationMetadata = new(content, messageId, false, $"Agent Framework Message {messageId}", correlationId) - { - SequenceNumber = DateTime.UtcNow.Ticks, - }; + long baseSequenceNumber = DateTime.UtcNow.Ticks; + List mappedContents = MapMessageContents(message); ActivityMetadata activityMetadata = new(activity); PolicyLocation policyLocation; @@ -158,20 +235,127 @@ private async Task> MapMessageToPCRequestsAsync(IEnu OperatingSystemVersion = "Unknown" } }; - ContentToProcess contentToProcess = new([conversationMetadata], activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata); if (string.IsNullOrEmpty(resolvedUserId)) { throw new PurviewRequestException("No user id provided or inferred for Purview request. Please provide an Entra user id in each message, pass a user id to the processor, or configure the TokenCredential to authenticate to an Entra user."); } - ProcessContentRequest pcRequest = new(contentToProcess, resolvedUserId, tenantId); - pcRequests.Add(pcRequest); + for (int entryIndex = 0; entryIndex < mappedContents.Count; entryIndex++) + { + string identifier = entryIndex == 0 ? messageId : $"{messageId}-{entryIndex}"; + ProcessConversationMetadata contentEntry = new(mappedContents[entryIndex], identifier, false, $"Agent Framework Message {messageId}", correlationId) + { + SequenceNumber = baseSequenceNumber + entryIndex, + }; + ContentToProcess contentToProcess = new(contentEntry, activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata); + ProcessContentRequest pcRequest = new(contentToProcess, resolvedUserId, tenantId); + pcRequests.Add(pcRequest); + } } return pcRequests; } + /// + /// Map every content item of a message onto the Purview content type that fits it. + /// + /// The message whose contents should be evaluated. + /// One content item per evaluable . + /// + /// and content with no data are skipped because they carry no user + /// data. Everything else is mapped to a real content item: submitting a message with part of + /// its payload unevaluated is a policy bypass. + /// + private static List MapMessageContents(ChatMessage message) + { + List mapped = []; + + foreach (AIContent content in message.Contents) + { + ContentBase? purviewContent = MapContent(content); + if (purviewContent != null) + { + mapped.Add(purviewContent); + } + } + + return mapped; + } + + /// + /// Map a single onto a Purview content item. + /// + /// The content item to map. + /// The Purview content item, or when the content carries no user data. + private static ContentBase? MapContent(AIContent content) + { + switch (content) + { + case UsageContent: + // Telemetry only; there is nothing for DLP to classify. + return null; + + case TextContent { Text: null or "" }: + case TextReasoningContent { Text: null or "" }: + case DataContent { Data.IsEmpty: true }: + // Empty data is skipped because the Graph APIs do not support it. A request with + // empty data cannot violate content policies because it contains no content. + return null; + + case TextContent textContent: + return new PurviewTextContent(textContent.Text); + + case TextReasoningContent reasoningContent: + return new PurviewTextContent(reasoningContent.Text); + + case DataContent dataContent: + return new PurviewBinaryContent(dataContent.Data.ToArray()); + + case UriContent uriContent: + return new PurviewTextContent($"{uriContent.MediaType} {uriContent.Uri}"); + + case FunctionCallContent functionCallContent: + return new PurviewTextContent(SerializeForEvaluation(functionCallContent)); + + case FunctionResultContent functionResultContent: + return new PurviewTextContent(SerializeForEvaluation(functionResultContent)); + + case ErrorContent errorContent: + return new PurviewTextContent(errorContent.Message ?? string.Empty); + + default: + // Catch-all for every remaining content type, including ones added after this code + // was written. Serializing is always safe; skipping would be a policy bypass. + return new PurviewTextContent(SerializeForEvaluation(content)); + } + } + + /// + /// Render a content item as text so Purview can classify everything it carries. + /// + /// The content item to render. + /// The text representation of . + /// The content could not be rendered for evaluation. + /// + /// The type info is resolved from the concrete runtime type rather than . + /// Polymorphic serialization through the base type throws for any subclass the abstractions do not + /// declare, which would push third-party and future content types onto the failure path. + /// + private static string SerializeForEvaluation(AIContent content) + { + try + { + return JsonSerializer.Serialize(content, AIJsonUtilities.DefaultOptions.GetTypeInfo(content.GetType())); + } + catch (Exception ex) when (ex is NotSupportedException or InvalidOperationException or JsonException) + { + // Fail closed. Falling back to ToString() would submit a bare type name in place of the + // payload, so the content would clear policy without ever having been evaluated. + throw new PurviewRequestException($"Unable to serialize content of type '{content.GetType()}' for Purview policy evaluation. Content cannot be evaluated and will not be released.", ex); + } + } + /// /// Orchestrates process content and protection scopes calls. /// @@ -228,12 +412,18 @@ private async Task ProcessWithCachedScopesAsync( if (shouldProcess) { - pcRequest.ProcessInline = executionMode == ExecutionMode.EvaluateInline; + pcRequest.ProcessInline = executionMode != ExecutionMode.EvaluateOffline; if (executionMode == ExecutionMode.EvaluateOffline) { this._channelHandler.QueueJob(new ProcessContentJob(pcRequest)); - return new ProcessContentResponse(); + + // Offline evaluation is asynchronous, but an explicit block action already known + // from the cached scopes must still be enforced rather than discarded. + List blockingActions = GetBlockingActions(dlpActions); + return blockingActions.Count > 0 + ? new ProcessContentResponse { PolicyActions = blockingActions } + : new ProcessContentResponse(); } return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions, cancellationToken).ConfigureAwait(false); @@ -327,7 +517,8 @@ internal static (bool shouldProcess, List dlpActions, ExecutionMo foreach (var location in scope.Locations ?? Array.Empty()) { - if (location.DataType.EndsWith(locationType, StringComparison.OrdinalIgnoreCase) && location.Value.Equals(locationValue, StringComparison.OrdinalIgnoreCase)) + if (location.DataType.EndsWith(locationType, StringComparison.OrdinalIgnoreCase) + && NormalizeLocationValue(locationType, location.Value).Equals(NormalizeLocationValue(locationType, locationValue), StringComparison.Ordinal)) { locationMatch = true; break; @@ -338,7 +529,9 @@ internal static (bool shouldProcess, List dlpActions, ExecutionMo { shouldProcess = true; - if (scope.ExecutionMode == ExecutionMode.EvaluateInline) + // Only an explicitly offline scope may skip inline evaluation. ExecutionMode is an + // evolvable enum, so any unrecognised value is upgraded to inline (fail closed). + if (scope.ExecutionMode != ExecutionMode.EvaluateOffline) { executionMode = ExecutionMode.EvaluateInline; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs index c3415d42af4..fa9c19878ca 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs @@ -519,7 +519,7 @@ private static ContentToProcess CreateValidContentToProcess() }; return new ContentToProcess( - [metadata], + metadata, activityMetadata, deviceMetadata, integratedAppMetadata, diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs index b59cc7a3929..7afad015ca1 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs @@ -296,6 +296,175 @@ public void CheckApplicableScopes_MatchesAnyLocationInScope() Assert.Equal(ExecutionMode.EvaluateInline, executionMode); } + [Fact] + public void CheckApplicableScopes_MatchesUrlLocationHostCaseInsensitively() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "HTTPS://Contoso.com/sites/marketing"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com/sites/marketing")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, _, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.True(shouldProcess); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationPathAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://contoso.com/sites/Marketing"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com/sites/marketing")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationQueryAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://contoso.com?label=Secret"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com?label=secret")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationFragmentAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://contoso.com#Section"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com#section")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationUserInfoAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://alice:SecretPass@contoso.com/docs"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://alice:secretpass@contoso.com/docs")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_MatchesApplicationLocationCaseInsensitively() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationApplication", "A1B2C3D4-E5F6-4A5B-8C9D-0E1F2A3B4C5D"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("#microsoft.graph.policyLocationApplication", "a1b2c3d4-e5f6-4a5b-8c9d-0e1f2a3b4c5d") + ], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, _, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.True(shouldProcess); + } + [Fact] public async Task ProcessMessagesAsync_UsesCachedProtectionScopes_WhenAvailableAsync() { @@ -993,10 +1162,287 @@ public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSub It.IsAny()), Times.Once); } + /// + /// Verifies every content item is submitted for evaluation in a separate request, not just message.Text. + /// Images, binary payloads and structured tool results are empty when flattened to text, which + /// would have them reach the model having only ever been classified as an empty string. + /// + [Fact] + public async Task ProcessMessagesAsync_WithMultipleContentItems_SubmitsEachInSeparateRequestAsync() + { + // Arrange + byte[] secret = [0x01, 0x02, 0x03, 0x04]; + var messages = new List + { + new(ChatRole.User, + [ + new DataContent(secret, "application/octet-stream"), + new FunctionCallContent("call-1", "exfiltrate", new Dictionary { ["ssn"] = "123-45-6789" }) + ]) + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); + + List capturedRequests = []; + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .Callback((request, _) => capturedRequests.Add(request)) + .ReturnsAsync(new ProcessContentResponse()); + + // Act + await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.Equal(2, capturedRequests.Count); + ProcessContentMetadataBase binaryEntry = Assert.Single(capturedRequests[0].ContentToProcess.ContentEntries); + ProcessContentMetadataBase functionCallEntry = Assert.Single(capturedRequests[1].ContentToProcess.ContentEntries); + + PurviewBinaryContent binaryContent = Assert.IsType(binaryEntry.Content); + Assert.Equal(secret, binaryContent.Data); + + PurviewTextContent functionCallContent = Assert.IsType(functionCallEntry.Content); + Assert.Contains("123-45-6789", functionCallContent.Data, StringComparison.Ordinal); + + // Content entries must be individually addressable, not collapsed onto one identifier. + Assert.NotEqual(binaryEntry.Identifier, functionCallEntry.Identifier); + } + + [Fact] + public async Task ProcessMessagesAsync_WithEmptyTextAndData_SkipsEmptyContentAsync() + { + // Arrange + byte[] emptyData = []; + List messages = + [ + new(ChatRole.User, + [ + new TextContent(string.Empty), + new DataContent(emptyData, "application/octet-stream") + ]) + ]; + PurviewSettings settings = CreateValidPurviewSettings(); + TokenInfo tokenInfo = new() { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + // Act + await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); + } + + /// + /// Verifies a block verdict known from an offline cached scope is still enforced. The offline + /// branch reports asynchronously, but discarding the scope's own policy actions would let a + /// known restrictAccess verdict go unenforced. + /// + [Fact] + public async Task ProcessMessagesAsync_WithOfflineScopeCarryingBlockAction_ReturnsShouldBlockTrueAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + var psResponse = new ProtectionScopesResponse + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ], + ExecutionMode = ExecutionMode.EvaluateOffline, + PolicyActions = + [ + new() { Action = DlpAction.RestrictAccess, RestrictionAction = RestrictionAction.Block } + ] + } + ] + }; + + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(psResponse); + + // Act + var result = await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.True(result.shouldBlock); + + // The offline report is still queued; enforcement is additional, not a replacement. + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + } + + /// + /// Verifies a restrictAccess scope whose restriction mode does not block is not enforced as + /// a block. The action carries a separate that may be + /// , or + /// , none of which withhold the content. + /// + [Fact] + public async Task ProcessMessagesAsync_WithOfflineScopeCarryingAuditRestriction_ReturnsShouldBlockFalseAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + var psResponse = new ProtectionScopesResponse + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ], + ExecutionMode = ExecutionMode.EvaluateOffline, + PolicyActions = + [ + new() { Action = DlpAction.RestrictAccess, RestrictionAction = RestrictionAction.Audit } + ] + } + ] + }; + + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(psResponse); + + // Act + var result = await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.False(result.shouldBlock); + + // The offline report is still queued so the audit action is recorded. + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + } + + /// + /// Verifies an unrecognised is evaluated inline. The enum is + /// evolvable, so an unknown member must not be treated as permission to skip enforcement. + /// + [Fact] + public async Task ProcessMessagesAsync_WithUnknownExecutionMode_EvaluatesInlineAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse((ExecutionMode)9999)); + + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(new ProcessContentResponse + { + PolicyActions = [new() { Action = DlpAction.BlockAccess }] + }); + + // Act + var result = await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.True(result.shouldBlock); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Once); + } + + /// + /// Verifies an subclass the abstractions do not declare still has its + /// payload evaluated. Serializing such a type through the polymorphic base type throws, and + /// substituting its type name would submit no payload at all for classification. + /// + [Fact] + public async Task ProcessMessagesAsync_WithUnknownContentSubclass_SubmitsPayloadForEvaluationAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, [new UnknownTestContent { Secret = "ssn 123-45-6789" }]) + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); + + ProcessContentRequest? capturedRequest = null; + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .Callback((request, _) => capturedRequest = request) + .ReturnsAsync(new ProcessContentResponse()); + + // Act + await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.NotNull(capturedRequest); + ProcessContentMetadataBase entry = Assert.Single(capturedRequest.ContentToProcess.ContentEntries); + PurviewTextContent content = Assert.IsType(entry.Content); + + Assert.Contains("123-45-6789", content.Data, StringComparison.Ordinal); + Assert.NotEqual(nameof(UnknownTestContent), content.Data); + } + #endregion #region Helper Methods + /// + /// Stands in for a third-party or future subclass that the abstractions do + /// not declare a polymorphic discriminator for. + /// + private sealed class UnknownTestContent : AIContent + { + public string Secret { get; set; } = string.Empty; + } + private static ProtectionScopesRequest CreateProtectionScopesRequest() { return new ProtectionScopesRequest("user-123", "tenant-123") @@ -1054,7 +1500,7 @@ private static ProcessContentRequest CreateProcessContentRequest() Version = "1.0" }; ContentToProcess contentToProcess = new( - [metadata], + metadata, activityMetadata, deviceMetadata, integratedAppMetadata, diff --git a/python/packages/purview/README.md b/python/packages/purview/README.md index 40a53b839a1..3b596a75c13 100644 --- a/python/packages/purview/README.md +++ b/python/packages/purview/README.md @@ -314,8 +314,46 @@ except (PurviewAuthenticationError, PurviewRateLimitError, PurviewRequestError, --- +## Security Considerations + +### Identity is a trusted input + +Purview evaluates DLP policy **for a specific user**. The identity this integration resolves therefore +decides *which* policy is applied, and it is resolved in this order: + +1. The `user_id` from the configured credential's token, when the credential resolves to a user. +2. The `user_id` argument passed to the processor. +3. `message.additional_properties["user_id"]`. +4. `message.author_name`, when it is a GUID. + +Only source 1 is verified. Sources 2–4 are supplied by the hosting application, so **a host must not +populate them from data that has crossed a trust boundary**. If an end user, an upstream service or a +model response can influence `additional_properties["user_id"]` or `author_name`, that party can +select a different user's DLP policy — typically one with weaker rules — and evade enforcement. Where +identity must come from a request, derive it from a validated token on the server, never from the +request body. Prefer a user-delegated credential (source 1) whenever possible. + +The `purview_app_location` in `PurviewSettings` is trusted in the same way: it selects which policy +locations apply and must be configured by the host, not by the caller. + +### Fail-closed behaviour + +Policy evaluation fails closed. If no user id can be resolved, or the tenant or app location cannot be +determined, the processor raises rather than letting content through unevaluated. Use +`ignore_exceptions` if you deliberately want availability over enforcement — but understand that it +disables enforcement for every error, not just transient ones. + +### What is evaluated + +Every content item on a message is submitted for evaluation, not just its text: binary/data content is +sent as Purview binary content, and function calls, function results and other structured content are +serialized to text. Only `usage` content is skipped, because it carries token counts rather than user +data. + +--- + ## Notes -- **User Identification**: When the configured credential resolves to a user token, that token's `user_id` is used for per-user policy scoping. For app-token credentials, provide a `user_id` per request (e.g. in `Message(..., additional_properties={"user_id": ""})`). If no user_id is provided or inferred, policy evaluation is skipped. +- **User Identification**: When the configured credential resolves to a user token, that token's `user_id` is used for per-user policy scoping. For app-token credentials, provide a `user_id` per request (e.g. in `Message(..., additional_properties={"user_id": ""})`). If no user_id can be provided or inferred, the request fails rather than proceeding unevaluated — see [Security Considerations](#security-considerations). - **Blocking Messages**: Can be customized via `blocked_prompt_message` and `blocked_response_message` in `PurviewSettings`. By default, they are "Prompt blocked by policy" and "Response blocked by policy" respectively. - **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. diff --git a/python/packages/purview/agent_framework_purview/_models.py b/python/packages/purview/agent_framework_purview/_models.py index 40df9d78cfe..76e1b88acc0 100644 --- a/python/packages/purview/agent_framework_purview/_models.py +++ b/python/packages/purview/agent_framework_purview/_models.py @@ -153,6 +153,7 @@ def serialize_flag( class DlpAction(str, Enum): BLOCK_ACCESS = "blockAccess" + RESTRICT_ACCESS = "restrictAccess" OTHER = "other" @@ -611,7 +612,7 @@ class ContentToProcess(_AliasSerializable): def __init__( self, - content_entries: list[ProcessConversationMetadata | MutableMapping[str, Any]], + content_entry: ProcessConversationMetadata | MutableMapping[str, Any], activity_metadata: ActivityMetadata | MutableMapping[str, Any], device_metadata: DeviceMetadata | MutableMapping[str, Any], integrated_app_metadata: IntegratedAppMetadata | MutableMapping[str, Any], @@ -619,8 +620,6 @@ def __init__( **kwargs: Any, ) -> None: # Extract aliased values from kwargs - if "contentEntries" in kwargs: - content_entries = kwargs["contentEntries"] if "activityMetadata" in kwargs: activity_metadata = kwargs["activityMetadata"] if "deviceMetadata" in kwargs: @@ -631,10 +630,11 @@ def __init__( protected_app_metadata = kwargs["protectedAppMetadata"] # Convert nested objects - entries = [ - e if isinstance(e, ProcessConversationMetadata) else ProcessConversationMetadata(**e) - for e in content_entries - ] + entry = ( + content_entry + if isinstance(content_entry, ProcessConversationMetadata) + else ProcessConversationMetadata(**content_entry) + ) if isinstance(activity_metadata, MutableMapping): activity_metadata = ActivityMetadata(**activity_metadata) if isinstance(device_metadata, MutableMapping): @@ -646,7 +646,7 @@ def __init__( # Call parent without explicit params with aliases super().__init__(**kwargs) - self.content_entries = entries + self.content_entries = [entry] self.activity_metadata = activity_metadata self.device_metadata = device_metadata self.integrated_app_metadata = integrated_app_metadata diff --git a/python/packages/purview/agent_framework_purview/_processor.py b/python/packages/purview/agent_framework_purview/_processor.py index d6c24a078f2..919e97723e3 100644 --- a/python/packages/purview/agent_framework_purview/_processor.py +++ b/python/packages/purview/agent_framework_purview/_processor.py @@ -1,13 +1,18 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import base64 +import binascii +import json import logging +import re import time import uuid from collections.abc import Iterable, MutableMapping +from copy import copy from typing import Any -from agent_framework import Message +from agent_framework import Content, Message from ._cache import CacheProvider, InMemoryCacheProvider, create_protection_scopes_cache_key from ._client import PurviewClient @@ -16,6 +21,7 @@ Activity, ActivityMetadata, ContentActivitiesRequest, + ContentBase, ContentToProcess, DeviceMetadata, DlpAction, @@ -31,6 +37,7 @@ ProtectionScopesRequest, ProtectionScopesResponse, ProtectionScopeState, + PurviewBinaryContent, PurviewTextContent, RestrictionAction, translate_activity, @@ -51,6 +58,126 @@ def _is_valid_guid(value: str | None) -> bool: return False +# RFC 2397: data:[][;base64],, where may be followed by any number of +# ";parameter=value" segments (for example "data:text/plain;charset=utf-8;base64,..."). The optional +# parameters have to be matched explicitly, otherwise a parameterised media type fails to decode and +# its payload would be submitted as a base64 string that no classifier can read. +_DATA_URI_PATTERN = re.compile( + r"^data:(?P[^;,]*)(?:;[^;,]+)*?;base64,(?P.*)$", + re.DOTALL | re.IGNORECASE, +) + +# Content types that carry no user data and therefore have nothing for DLP to classify. +# Every other content type must reach Purview; see _map_content. +_NON_EVALUATED_CONTENT_TYPES = frozenset({"usage"}) + + +def _serialize_for_evaluation(value: Any) -> str: + """Render an arbitrary content value as text so Purview can classify it.""" + if value is None: + return "" + if isinstance(value, str): + return value + try: + return json.dumps(value, default=str, sort_keys=True) + except (TypeError, ValueError): + return str(value) + + +def _decode_data_uri(uri: str | None) -> bytes | None: + """Return the raw bytes behind a base64 data URI, or None if it is not one.""" + if not uri: + return None + match = _DATA_URI_PATTERN.match(uri) + if not match: + return None + try: + return base64.b64decode(match.group("base64_data"), validate=True) + except (binascii.Error, ValueError): + return None + + +def _map_content(content: Content) -> ContentBase | None: + """Map a single message content item onto the Purview content type that fits it. + + Returns None only for content that carries no user data at all. Anything else is + mapped to a real content entry: a message must never reach the model with part of + its payload unevaluated. + """ + content_type = content.type + + if content_type in _NON_EVALUATED_CONTENT_TYPES: + return None + + # Empty data is skipped because the Graph APIs do not support it. A request with + # empty data cannot violate content policies because it contains no content. + if content_type in ("text", "text_reasoning"): + if not content.text: + return None + return PurviewTextContent(data=content.text) + + if content_type == "data": + uri = getattr(content, "uri", None) + if not uri: + return None + raw_data = _decode_data_uri(uri) + if raw_data is not None: + if not raw_data: + return None + return PurviewBinaryContent(data=raw_data) + # Not a base64 data URI after all: evaluate the serialized form rather than drop it. + return PurviewTextContent(data=_serialize_for_evaluation(content.to_dict())) + + # Everything else - uri, function_call, function_result and any content type added after this + # code was written - is serialized whole. Serializing the entire item rather than picking out + # named fields keeps additional_properties, which is a data channel in its own right, under + # evaluation. Skipping a content type would be a policy bypass. + return PurviewTextContent(data=_serialize_for_evaluation(content.to_dict())) + + +def _map_message_contents(message: Message) -> list[ContentBase]: + """Map every non-empty content item of a message to Purview content entries.""" + return [purview_content for content in message.contents if (purview_content := _map_content(content))] + + +def _is_blocking_action(action_info: DlpActionInfo) -> bool: + """Whether a policy action means the content must not be released. + + ``restrictAccess`` is not blocking on its own: it carries a separate ``restrictionAction`` that + selects the enforcement mode, which may be audit, warn or allow as well as block. Only an explicit + block mode withholds the content. + """ + return action_info.action == DlpAction.BLOCK_ACCESS or action_info.restriction_action == RestrictionAction.BLOCK + + +def _blocking_actions(dlp_actions: list[DlpActionInfo]) -> list[DlpActionInfo | MutableMapping[str, Any]]: + """Filter policy actions down to the ones that block.""" + return [action_info for action_info in dlp_actions if _is_blocking_action(action_info)] + + +def _normalize_location_value(data_type: str, value: str) -> str: + """Normalize a policy location value for comparison, according to its location type. + + Application ids (GUIDs) and domain names are case-insensitive, so those fold whole. URL values are + not. Only the scheme and the host are case-insensitive; the userinfo, path, query and fragment are + all case-sensitive, so folding a URL whole would let a scope for ``contoso.com/public`` match a + request for ``contoso.com/Public``. Location types that are not recognised fold whole, which matches + more scopes rather than fewer. + """ + if data_type.split(".")[-1].casefold().endswith("url"): + scheme, separator, remainder = value.partition("://") + if not separator: + scheme, separator, remainder = "", "", value + # The authority ends at the first path, query or fragment delimiter, whichever comes first. + delimiter = re.search(r"[/?#]", remainder) + authority_end = delimiter.start() if delimiter else len(remainder) + authority, tail = remainder[:authority_end], remainder[authority_end:] + # Credentials are case-sensitive, so only the host half of the authority folds. + userinfo, at_sign, host = authority.rpartition("@") + return f"{scheme.casefold()}{separator}{userinfo}{at_sign}{host.casefold()}{tail}" + return value.casefold() + + class ScopedContentProcessor: """Combine protection scopes, process content, and content activities logic.""" @@ -91,7 +218,7 @@ async def process_messages( resp = await self._process_with_scopes(req) if resp.policy_actions: for act in resp.policy_actions: - if act.action == DlpAction.BLOCK_ACCESS or act.restriction_action == RestrictionAction.BLOCK: + if _is_blocking_action(act): should_block = True break if should_block: @@ -141,23 +268,21 @@ async def _map_messages( if not resolved_user_id and resolved_author_name: resolved_user_id = resolved_author_name - # Return empty results if user_id is empty + # Fail closed: without a resolvable user identity no policy can be evaluated, + # so the content must not be allowed through unevaluated. if not resolved_user_id or not _is_valid_guid(resolved_user_id): - return results, None + raise ValueError( + "No user id provided or inferred for Purview request. Please provide an Entra user id in each " + "message, pass a user id to the processor, or configure the credential to authenticate to an " + "Entra user." + ) for m in messages: message_id = m.message_id or str(uuid.uuid4()) - content = PurviewTextContent(data=m.text or "") correlation_id = (session_id or str(uuid.uuid4())) + "@AF" - meta = ProcessConversationMetadata( - identifier=message_id, - content=content, - name=f"Agent Framework Message {message_id}", - is_truncated=False, - correlation_id=correlation_id, - # This would be c# ticks equivalent and needs to fit inside c# long - sequence_number=time.time_ns() // 100 + 621355968000000000, - ) + # This would be c# ticks equivalent and needs to fit inside c# long + base_sequence_number = time.time_ns() // 100 + 621355968000000000 + mapped_contents = _map_message_contents(m) activity_meta = ActivityMetadata(activity=activity) purview_app_location = self._settings.get("purview_app_location") @@ -187,21 +312,30 @@ async def _map_messages( ) ) - ctp = ContentToProcess( - content_entries=[meta], - activity_metadata=activity_meta, - device_metadata=device_meta, - integrated_app_metadata=integrated_app, - protected_app_metadata=protected_app, - ) - req = ProcessContentRequest( - content_to_process=ctp, - user_id=resolved_user_id, # Use the resolved user_id for all messages - tenant_id=tenant_id, - correlation_id=meta.correlation_id, - process_inline=None, # Will be set based on execution mode - ) - results.append(req) + for index, purview_content in enumerate(mapped_contents): + content_entry = ProcessConversationMetadata( + identifier=message_id if index == 0 else f"{message_id}-{index}", + content=purview_content, + name=f"Agent Framework Message {message_id}", + is_truncated=False, + correlation_id=correlation_id, + sequence_number=base_sequence_number + index, + ) + ctp = ContentToProcess( + content_entry=content_entry, + activity_metadata=activity_meta, + device_metadata=device_meta, + integrated_app_metadata=integrated_app, + protected_app_metadata=protected_app, + ) + req = ProcessContentRequest( + content_to_process=ctp, + user_id=resolved_user_id, # Use the resolved user_id for all messages + tenant_id=tenant_id, + correlation_id=correlation_id, + process_inline=None, # Will be set based on execution mode + ) + results.append(req) return results, resolved_user_id async def _process_with_scopes(self, pc_request: ProcessContentRequest) -> ProcessContentResponse: @@ -231,7 +365,9 @@ async def _process_with_scopes(self, pc_request: ProcessContentRequest) -> Proce return await self._process_with_cached_scopes(pc_request, cached_ps_resp, cache_key) pc_request.process_inline = True - task = asyncio.create_task(self._refresh_protection_scopes_background(ps_req, cache_key, pc_request)) + # The background refresh gets its own copy: the foreground call mutates + # process_inline and scope_identifier on this request while that task runs. + task = asyncio.create_task(self._refresh_protection_scopes_background(ps_req, cache_key, copy(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=[]) @@ -248,15 +384,23 @@ async def _process_with_cached_scopes( should_process, dlp_actions, execution_mode = self._check_applicable_scopes(pc_request, ps_resp) if should_process: - # Set process_inline based on execution mode - pc_request.process_inline = execution_mode == ExecutionMode.EVALUATE_INLINE + # Only an explicitly offline scope may skip inline evaluation. executionMode is an + # evolvable enum, so any unrecognised value must fail closed and be evaluated inline. + evaluate_offline = execution_mode == ExecutionMode.EVALUATE_OFFLINE + pc_request.process_inline = not evaluate_offline # If execution mode is offline, queue the PC request in background - if execution_mode != ExecutionMode.EVALUATE_INLINE: + if evaluate_offline: task = asyncio.create_task(self._process_content_background(pc_request, cache_key)) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) - return ProcessContentResponse(id="204", correlation_id=pc_request.correlation_id) + # Offline evaluation is asynchronous, but an explicit block action already known + # from the cached scopes must still be enforced rather than discarded. + return ProcessContentResponse( + id="204", + correlation_id=pc_request.correlation_id, + policy_actions=_blocking_actions(dlp_actions) or None, + ) return await self._call_process_content(pc_request, cache_key, dlp_actions=dlp_actions) @@ -373,15 +517,19 @@ def _check_applicable_scopes( loc.data_type and location.data_type and loc.data_type.lower().endswith(location.data_type.split(".")[-1].lower()) - and loc.value == location.value + and isinstance(loc.value, str) + and isinstance(location.value, str) + and _normalize_location_value(location.data_type, loc.value) + == _normalize_location_value(location.data_type, location.value) ): location_match = True break if activity_match and location_match: should_process = True - # If any scope has EvaluateInline, upgrade to inline mode - if scope.execution_mode == ExecutionMode.EVALUATE_INLINE: + # Only an explicitly offline scope may skip inline evaluation. execution_mode is an + # evolvable enum, so any unrecognised value is upgraded to inline (fail closed). + if scope.execution_mode != ExecutionMode.EVALUATE_OFFLINE: execution_mode = ExecutionMode.EVALUATE_INLINE if scope.policy_actions: diff --git a/python/packages/purview/tests/purview/conftest.py b/python/packages/purview/tests/purview/conftest.py index dc9a7024e8c..0089053b7c9 100644 --- a/python/packages/purview/tests/purview/conftest.py +++ b/python/packages/purview/tests/purview/conftest.py @@ -41,7 +41,7 @@ def _create_content(text: str = "Test") -> ContentToProcess: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) return ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, diff --git a/python/packages/purview/tests/purview/test_processor.py b/python/packages/purview/tests/purview/test_processor.py index 9d7388a416b..f5d2fe81efb 100644 --- a/python/packages/purview/tests/purview/test_processor.py +++ b/python/packages/purview/tests/purview/test_processor.py @@ -20,6 +20,8 @@ ProcessContentResponse, ProtectionScopeActivities, ProtectionScopeState, + PurviewBinaryContent, + PurviewTextContent, RestrictionAction, ) from agent_framework_purview._processor import ScopedContentProcessor, _is_valid_guid @@ -144,6 +146,126 @@ async def test_map_messages_creates_requests( assert requests[0].tenant_id == "12345678-1234-1234-1234-123456789012" assert user_id == "12345678-1234-1234-1234-123456789012" + async def test_map_messages_submits_each_content_item_in_separate_request( + self, processor: ScopedContentProcessor + ) -> None: + """Test _map_messages submits every content item in a separate request. + + Non-text content flattened to Message.text is empty for binary, tool-call and + tool-result content, which would have Purview classify an empty string. + """ + from agent_framework import Content + + secret = b"credit card 4532667785213500" + messages = [ + Message( + role="user", + contents=[ + Content.from_data(data=secret, media_type="application/octet-stream"), + Content( + "function_call", + call_id="call-1", + name="send_email", + arguments={"body": "ssn 120-98-1437"}, + ), + Content("function_result", call_id="call-1", result="account 999-12345"), + ], + ) + ] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert len(requests) == 3 + entries = [request.content_to_process.content_entries[0] for request in requests] + assert all(len(request.content_to_process.content_entries) == 1 for request in requests) + binary_entry, call_entry, result_entry = entries + assert isinstance(binary_entry.content, PurviewBinaryContent) + assert binary_entry.content.data == secret + + assert isinstance(call_entry.content, PurviewTextContent) + assert "send_email" in call_entry.content.data + assert "120-98-1437" in call_entry.content.data + + assert isinstance(result_entry.content, PurviewTextContent) + assert "999-12345" in result_entry.content.data + + # Every entry must carry real content; an empty payload would not be evaluated. + assert all(entry.content is not None for entry in entries) + assert not any(isinstance(entry.content, PurviewTextContent) and entry.content.data == "" for entry in entries) + + async def test_map_messages_skips_empty_text_and_data(self, processor: ScopedContentProcessor) -> None: + """Test _map_messages does not create requests for empty text or binary data.""" + from agent_framework import Content + + messages = [ + Message( + role="user", + contents=[ + "", + Content.from_data(data=b"", media_type="application/octet-stream"), + ], + ) + ] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert requests == [] + + async def test_map_messages_decodes_data_uri_with_media_type_parameters( + self, processor: ScopedContentProcessor + ) -> None: + """Test _map_messages decodes base64 data URIs whose media type carries parameters. + + A data URI may carry any number of ";parameter=value" segments between the media type and + ";base64". Failing to decode one leaves the payload as a base64 string that no classifier + can read. + """ + from agent_framework import Content + + secret = b"credit card 4532667785213500" + content = Content.from_data(data=secret, media_type="text/plain;charset=utf-8") + assert content.uri is not None + assert content.uri.startswith("data:text/plain;charset=utf-8;base64,") + + messages = [Message(role="user", contents=[content])] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + entries = requests[0].content_to_process.content_entries + assert len(entries) == 1 + assert isinstance(entries[0].content, PurviewBinaryContent) + assert entries[0].content.data == secret + + async def test_map_messages_submits_additional_properties_on_structured_content( + self, processor: ScopedContentProcessor + ) -> None: + """Test _map_messages serializes structured content whole, including additional_properties. + + Building an entry from name/arguments alone drops additional_properties, leaving a + data channel on tool calls and tool results that Purview never sees. + """ + from agent_framework import Content + + call = Content("function_call", call_id="call-1", name="send_email", arguments={"body": "hello"}) + call.additional_properties = {"hidden": "ssn 120-98-1437"} + result = Content("function_result", call_id="call-1", result="ok") + result.additional_properties = {"hidden": "card 4532667785213500"} + + messages = [Message(role="user", contents=[call, result])] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert len(requests) == 2 + call_entry = requests[0].content_to_process.content_entries[0] + result_entry = requests[1].content_to_process.content_entries[0] + + assert isinstance(call_entry.content, PurviewTextContent) + assert "120-98-1437" in call_entry.content.data + assert "send_email" in call_entry.content.data + + assert isinstance(result_entry.content, PurviewTextContent) + assert "4532667785213500" in result_entry.content.data + async def test_map_messages_without_defaults_gets_token_info(self, mock_client: AsyncMock) -> None: """Test _map_messages gets token info when settings lack some defaults.""" settings = PurviewSettings(app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012") @@ -220,6 +342,207 @@ async def test_check_applicable_scopes_with_block_action( assert len(actions) == 1 assert actions[0].action == DlpAction.BLOCK_ACCESS + async def test_check_applicable_scopes_matches_location_case_insensitively( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes matches location values regardless of GUID casing. + + A casing difference between the request location and the scope location would + otherwise hide an applicable block scope. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationApplication", + value="A1B2C3D4-E5F6-4A5B-8C9D-0E1F2A3B4C5D", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationApplication", + value="a1b2c3d4-e5f6-4a5b-8c9d-0e1f2a3b4c5d", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, execution_mode = ScopedContentProcessor._check_applicable_scopes( + pc_request, ps_response + ) + + assert should_process is True + assert execution_mode == ExecutionMode.EVALUATE_INLINE + assert dlp_actions + + async def test_check_applicable_scopes_matches_url_location_host_case_insensitively( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes ignores host casing on URL locations. + + The scheme and host of a URL are case-insensitive, so a casing difference there must not + hide an applicable scope. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="HTTPS://Contoso.com/sites/marketing", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com/sites/marketing", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is True + assert dlp_actions + + async def test_check_applicable_scopes_treats_url_location_path_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL path casing significant. + + URL paths are case-sensitive, so a scope scoped to one path must not match a different path + that differs only in casing. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://contoso.com/sites/Marketing", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com/sites/marketing", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + + async def test_check_applicable_scopes_treats_url_location_query_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL query casing significant. + + A query value is case-sensitive, and it may follow the host directly with no path between + them, so the authority has to end at the query delimiter as well as the path delimiter. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://contoso.com?label=Secret", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com?label=secret", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + + async def test_check_applicable_scopes_treats_url_location_fragment_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL fragment casing significant. + + A fragment may follow the host directly, so the authority has to end at the fragment + delimiter as well. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://contoso.com#Section", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com#section", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + + async def test_check_applicable_scopes_treats_url_location_userinfo_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL userinfo casing significant. + + Credentials embedded in the authority are case-sensitive, so only the host half of the + authority may be folded. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://alice:SecretPass@contoso.com/docs", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://alice:secretpass@contoso.com/docs", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + async def test_combine_policy_actions(self, processor: ScopedContentProcessor) -> None: """Test _combine_policy_actions merges action lists.""" action1 = DlpActionInfo(action=DlpAction.BLOCK_ACCESS, restrictionAction=RestrictionAction.BLOCK) @@ -290,6 +613,115 @@ async def test_process_with_scopes_calls_client_methods( mock_client.get_protection_scopes.assert_called_once() mock_client.send_content_activities.assert_called_once() + async def test_process_with_scopes_isolates_the_background_refresh_request( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test a cold scopes cache evaluates inline and hands the background refresh its own request. + + The background refresh must not observe foreground mutations of the request it was given. + """ + pc_request = process_content_request_factory() + cast(Any, processor._cache).get = AsyncMock(return_value=None) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + captured: list[Any] = [] + + async def capture_refresh(ps_req: Any, cache_key: str, request: Any) -> None: + captured.append(request) + + cast(Any, processor)._refresh_protection_scopes_background = capture_refresh + + await processor._process_with_scopes(pc_request) + await asyncio.gather(*list(processor._background_tasks)) + + # Content was evaluated inline, so there is no window in which it goes unevaluated. + mock_client.process_content.assert_called_once() + assert pc_request.process_inline is True + + # The background task got its own request object. + assert captured and captured[0] is not pc_request + + async def test_process_with_scopes_evaluates_unknown_execution_mode_inline( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test an unrecognised executionMode is evaluated inline. + + executionMode is an evolvable enum, so an unknown member must not skip enforcement. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.UNKNOWN_FUTURE_VALUE, + locations=[PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")], + policy_actions=[], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + cast(Any, processor._cache).get = AsyncMock(side_effect=[None, ps_response]) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + await processor._process_with_scopes(pc_request) + + mock_client.process_content.assert_called_once() + assert pc_request.process_inline is True + + async def test_process_with_scopes_reports_block_action_on_offline_scope( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test a block verdict known from an offline scope is still reported rather than discarded.""" + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_OFFLINE, + locations=[PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")], + policy_actions=[ + DlpActionInfo(action=DlpAction.RESTRICT_ACCESS, restriction_action=RestrictionAction.BLOCK) + ], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + cast(Any, processor._cache).get = AsyncMock(side_effect=[None, ps_response]) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + response = await processor._process_with_scopes(pc_request) + await asyncio.gather(*list(processor._background_tasks)) + + assert response.policy_actions + assert all( + action.action == DlpAction.RESTRICT_ACCESS and action.restriction_action == RestrictionAction.BLOCK + for action in response.policy_actions + ) + + async def test_process_with_scopes_ignores_non_blocking_restriction_on_offline_scope( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test a restrictAccess scope whose restriction mode does not block is not reported as blocking. + + restrictAccess carries a separate restriction mode which may be audit, warn or allow. Only an + explicit block mode withholds the content, so the other modes must not surface as a verdict. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_OFFLINE, + locations=[PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")], + policy_actions=[ + DlpActionInfo(action=DlpAction.RESTRICT_ACCESS, restriction_action=RestrictionAction.OTHER) + ], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + cast(Any, processor._cache).get = AsyncMock(side_effect=[None, ps_response]) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + response = await processor._process_with_scopes(pc_request) + await asyncio.gather(*list(processor._background_tasks)) + + assert not response.policy_actions + async def test_process_with_scopes_preserves_restriction_only_policy_actions( self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory ) -> None: @@ -516,8 +948,8 @@ async def test_map_messages_with_provided_user_id_fallback(self, mock_client: As assert user_id == "32345678-1234-1234-1234-123456789012" assert requests[0].user_id == "32345678-1234-1234-1234-123456789012" - async def test_map_messages_returns_empty_when_no_user_id(self, mock_client: AsyncMock) -> None: - """Test that empty results are returned when user_id cannot be resolved.""" + async def test_map_messages_raises_when_no_user_id(self, mock_client: AsyncMock) -> None: + """Test that an unresolvable user_id fails closed instead of skipping evaluation.""" settings = PurviewSettings( app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012", @@ -533,10 +965,8 @@ async def test_map_messages_returns_empty_when_no_user_id(self, mock_client: Asy messages = [Message(role="user", contents=["Test message"])] - requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) - - assert len(requests) == 0 - assert user_id is None + with pytest.raises(ValueError, match="No user id"): + await processor._map_messages(messages, Activity.UPLOAD_TEXT) async def test_process_content_sends_activities_when_not_applicable( self, mock_client: AsyncMock, process_content_request_factory @@ -705,11 +1135,9 @@ async def test_author_name_ignored_if_not_valid_guid( ) ] - requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) - - # Should return empty since author_name is not a valid GUID - assert user_id is None - assert len(requests) == 0 + # author_name is not a valid GUID, so no identity resolves and evaluation fails closed + with pytest.raises(ValueError, match="No user id"): + await processor._map_messages(messages, Activity.UPLOAD_TEXT) async def test_provided_user_id_used_as_last_resort( self, mock_client: AsyncMock, settings: PurviewSettings @@ -739,10 +1167,29 @@ async def test_invalid_provided_user_id_ignored(self, mock_client: AsyncMock, se messages = [Message(role="user", contents=["Test"])] - requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT, provided_user_id="not-a-guid") + with pytest.raises(ValueError, match="No user id"): + await processor._map_messages(messages, Activity.UPLOAD_TEXT, provided_user_id="not-a-guid") - assert user_id is None - assert len(requests) == 0 + async def test_unresolvable_user_id_blocks_processing( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test process_messages raises rather than sending content when no user id resolves. + + Without a user id there is no policy to evaluate against, so the content must not be + forwarded to Purview or on to the model. + """ + mock_client.get_user_info_from_token.return_value = { + "tenant_id": "12345678-1234-1234-1234-123456789012", + "client_id": "12345678-1234-1234-1234-123456789012", + } + processor = ScopedContentProcessor(mock_client, settings) + + messages = [Message(role="user", contents=["ssn 120-98-1437"])] + + with pytest.raises(ValueError, match="No user id"): + await processor.process_messages(messages, Activity.UPLOAD_TEXT) + + mock_client.process_content.assert_not_called() async def test_multiple_messages_same_user_id(self, mock_client: AsyncMock, settings: PurviewSettings) -> None: """Test that all messages use the same resolved user_id.""" diff --git a/python/packages/purview/tests/purview/test_purview_models.py b/python/packages/purview/tests/purview/test_purview_models.py index 4af581014d9..3e3e3fb929a 100644 --- a/python/packages/purview/tests/purview/test_purview_models.py +++ b/python/packages/purview/tests/purview/test_purview_models.py @@ -91,7 +91,7 @@ def test_content_to_process_with_nested_structures(self) -> None: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) content = ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, @@ -188,7 +188,7 @@ def test_content_serialization_uses_aliases(self) -> None: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) content = ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, @@ -198,6 +198,7 @@ def test_content_serialization_uses_aliases(self) -> None: dumped = content.model_dump(by_alias=True, exclude_none=True, mode="json") assert "contentEntries" in dumped + assert len(dumped["contentEntries"]) == 1 assert "activityMetadata" in dumped assert "deviceMetadata" in dumped assert "integratedAppMetadata" in dumped @@ -224,7 +225,7 @@ def test_process_content_request_excludes_private_fields(self) -> None: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) content = ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, From 0c07eb6a3de32cf1afcda37263294b2f8e07947c Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Thu, 17 Sep 2026 11:34:13 +0000 Subject: [PATCH 06/26] Python: Add per-tool AgentModeProvider exposure controls (#8450) * Python: Add mode tool exposure controls * Python: Support replacement mode state updates * Python: Clear stale mode notifications --- .../core/agent_framework/_harness/_mode.py | 86 ++++++++-- .../core/tests/core/test_harness_agent.py | 54 +++++++ .../core/tests/core/test_harness_mode.py | 151 ++++++++++++++++++ .../02-agents/context_providers/README.md | 27 ++++ 4 files changed, 301 insertions(+), 17 deletions(-) diff --git a/python/packages/core/agent_framework/_harness/_mode.py b/python/packages/core/agent_framework/_harness/_mode.py index 11cc018b178..9d8e1650286 100644 --- a/python/packages/core/agent_framework/_harness/_mode.py +++ b/python/packages/core/agent_framework/_harness/_mode.py @@ -8,18 +8,26 @@ from .._sessions import AgentSession, ContextProvider, SessionContext from .._telemetry import FeatureIndex, mark_feature_used -from .._tools import tool +from .._tools import FunctionTool, tool from .._types import Message DEFAULT_MODE_SOURCE_ID = "agent_mode" +_MODE_GET_INSTRUCTIONS = "Use the mode_get tool to check your current operating mode.\n" +_MODE_SET_INSTRUCTIONS = ( + "Use the mode_set tool to switch between modes as your work progresses. " + "Only use mode_set if the user explicitly instructs/allows you to change modes.\n\n" +) +_PLAN_MODE_TRANSITION = ( + "7. When approval is granted, always switch to execute mode (using the `mode_set` tool), " + "and follow the steps for *Execute mode*." +) DEFAULT_MODE_INSTRUCTIONS = ( "## Agent Mode\n\n" "- You can operate in different modes. Depending on the mode you are in, " "you will be required to follow different processes.\n\n" - "Use the mode_get tool to check your current operating mode.\n" - "Use the mode_set tool to switch between modes as your work progresses. " - "Only use mode_set if the user explicitly instructs/allows you to change modes.\n\n" - "You are currently operating in the {current_mode} mode.\n\n" + + _MODE_GET_INSTRUCTIONS + + _MODE_SET_INSTRUCTIONS + + "You are currently operating in the {current_mode} mode.\n\n" "### Mandatory Mode based Workflow\n\n" "For every new substantive user request, including short factual questions, " "your behavior is determined by the mode you are in.\n\n" @@ -49,8 +57,7 @@ "5. Write the plan to a memory file, so that it is retained even if compaction happens. " "Make sure to update the plan file if the user requests changes.\n" "6. Present the plan to the user and ask for approval to switch to execute mode and process the plan.\n" - "7. When approval is granted, always switch to execute mode (using the `mode_set` tool), " - "and follow the steps for *Execute mode*." + + _PLAN_MODE_TRANSITION ), "execute": ( "Determine the type of ask:\n" @@ -157,6 +164,7 @@ def set_agent_mode( *, source_id: str = DEFAULT_MODE_SOURCE_ID, available_modes: Sequence[str] | None = None, + notify: bool = True, ) -> str: """Set the current operating mode in session state. @@ -173,6 +181,9 @@ def set_agent_mode( Keyword Args: source_id: Unique source ID for the provider state. available_modes: Supported modes to validate against. Defaults to the built-in modes. + notify: Whether to notify the agent about the mode change on its next run. Set to ``False`` when the + agent changes mode through a replacement tool and has already observed the tool result. This also + clears any pending external-change notification. Returns: The normalized mode string that was stored. @@ -189,8 +200,11 @@ def set_agent_mode( # prior mode so the next ``before_run`` can inject a user message announcing the switch. Without # that injection, the model often anchors on the earlier ``set_mode`` tool call in the chat history # and keeps behaving as if it were still in that mode — system instructions alone are insufficient. - if isinstance(previous_mode, str) and previous_mode != normalized_mode: - provider_state[_PREVIOUS_MODE_STATE_KEY] = previous_mode + if notify: + if isinstance(previous_mode, str) and previous_mode != normalized_mode: + provider_state[_PREVIOUS_MODE_STATE_KEY] = previous_mode + else: + provider_state.pop(_PREVIOUS_MODE_STATE_KEY, None) return normalized_mode @@ -204,10 +218,13 @@ class AgentModeProvider(ContextProvider): The set of available modes is configurable with ``mode_instructions``. By default, two modes are provided: ``"plan"`` (interactive planning) and ``"execute"`` (autonomous execution). - This provider exposes the following tools to the agent: + By default, this provider exposes the following tools to the agent: - ``mode_set``: Switch the agent's operating mode. - ``mode_get``: Retrieve the agent's current operating mode. + Set ``expose_mode_set`` or ``expose_mode_get`` to ``False`` to omit that tool while retaining mode state + and workflow instructions. Replacement tools can be supplied through the agent's ``tools`` argument. + Public helper functions ``get_agent_mode`` and ``set_agent_mode`` allow external code to programmatically read and change the mode. """ @@ -219,6 +236,8 @@ def __init__( default_mode: str | None = None, mode_instructions: Mapping[str, str] | None = None, instructions: str | None = None, + expose_mode_set: bool = True, + expose_mode_get: bool = True, ) -> None: """Initialize a new agent mode provider. @@ -229,15 +248,30 @@ def __init__( default_mode: Initial mode used when no mode is stored yet. When omitted, the first entry of ``mode_instructions`` is used. mode_instructions: Mapping of supported modes to instructions on when and how to use each mode. + Custom text is not rewritten when tools are hidden. instructions: Custom instructions for using the mode tools. The instructions can contain an ``{available_modes}`` placeholder for the configured list of modes and a ``{current_mode}`` placeholder for the currently active mode. When omitted, the provider uses a default set of instructions. + Default guidance reflects tool exposure; custom text is not rewritten when tools are hidden. + expose_mode_set: Whether to contribute the built-in ``mode_set`` tool. Defaults to ``True``. + When ``False``, the application controls mode changes, optionally through a replacement tool. + expose_mode_get: Whether to contribute the built-in ``mode_get`` tool. Defaults to ``True``. + The current mode remains available in the default instructions and through ``get_agent_mode``. Raises: ValueError: No modes are configured, or the default mode is not configured. """ super().__init__(source_id) - mode_instructions = dict(DEFAULT_MODE_MAP if mode_instructions is None else mode_instructions) + if mode_instructions is None: + mode_instructions = dict(DEFAULT_MODE_MAP) + if not expose_mode_set: + mode_instructions["plan"] = mode_instructions["plan"].replace( + _PLAN_MODE_TRANSITION, + "7. When approval is granted, use the application's configured mode-change mechanism to " + "transition to execute mode. Follow the steps for *Execute mode* only after the mode has changed.", + ) + else: + mode_instructions = dict(mode_instructions) self._mode_display_names = _normalize_available_modes(tuple(mode_instructions)) if not self._mode_display_names: raise ValueError("mode_instructions must contain at least one mode.") @@ -247,6 +281,8 @@ def __init__( self.available_modes = tuple(self._mode_display_names) self.default_mode = _resolve_default_mode(default_mode, available_modes=self._mode_display_names) self.instructions = instructions + self.expose_mode_set = expose_mode_set + self.expose_mode_get = expose_mode_get def _build_instructions(self, current_mode: str) -> str: """Build the mode guidance injected for the current session.""" @@ -255,6 +291,15 @@ def _build_instructions(self, current_mode: str) -> str: for mode, mode_instruction in self.mode_instructions.items() ) instructions = self.instructions or DEFAULT_MODE_INSTRUCTIONS + if not self.instructions: + if not self.expose_mode_get: + instructions = instructions.replace(_MODE_GET_INSTRUCTIONS, "") + if not self.expose_mode_set: + instructions = instructions.replace( + _MODE_SET_INSTRUCTIONS, + "Mode changes are controlled by the application. Use its configured mode-change mechanism " + "only when the user explicitly instructs/allows a mode change.\n\n", + ) return instructions.replace("{available_modes}", mode_lines).replace("{current_mode}", current_mode) async def before_run( @@ -289,11 +334,13 @@ async def before_run( @tool(name="mode_set", approval_mode="never_require") def mode_set(mode: str) -> str: """Switch the agent's operating mode.""" - # The agent invoked the tool itself, so it knows the mode just changed — bypass - # ``set_agent_mode`` to avoid triggering a notification message on the next turn. - normalized_mode = _normalize_mode(mode, available_modes=self._mode_display_names) - tool_state = _get_mode_state(session, source_id=self.source_id) - tool_state["current_mode"] = normalized_mode + normalized_mode = set_agent_mode( + session, + mode, + source_id=self.source_id, + available_modes=self.available_modes, + notify=False, + ) return json.dumps({"mode": normalized_mode, "message": f"Mode changed to '{normalized_mode}'."}) @tool(name="mode_get", approval_mode="never_require") @@ -311,7 +358,12 @@ def mode_get() -> str: self.source_id, [self._build_instructions(current_mode)], ) - context.extend_tools(self.source_id, [mode_set, mode_get]) + tools: list[FunctionTool] = [] + if self.expose_mode_set: + tools.append(mode_set) + if self.expose_mode_get: + tools.append(mode_get) + context.extend_tools(self.source_id, tools) if isinstance(previous_mode, str) and previous_mode != current_mode: # Inject a user-role message announcing the external mode change. System instructions # always render first in the chat history, so the agent can otherwise stay anchored to diff --git a/python/packages/core/tests/core/test_harness_agent.py b/python/packages/core/tests/core/test_harness_agent.py index 6535a6fc68f..41030e3bd08 100644 --- a/python/packages/core/tests/core/test_harness_agent.py +++ b/python/packages/core/tests/core/test_harness_agent.py @@ -24,6 +24,7 @@ FileAccessProvider, FileMemoryProvider, FileSystemAgentFileStore, + FunctionTool, InMemoryAgentFileStore, InMemoryHistoryProvider, Message, @@ -33,6 +34,9 @@ SkillsProvider, TodoProvider, create_harness_agent, + get_agent_mode, + set_agent_mode, + tool, ) from agent_framework._harness._agent import DEFAULT_HARNESS_INSTRUCTIONS, _assemble_instructions from agent_framework._harness._mode import AgentModeProvider @@ -133,6 +137,56 @@ def test_create_harness_agent_disable_mode() -> None: assert AgentModeProvider not in provider_types +async def test_create_harness_agent_with_replacement_mode_tool() -> None: + """Applications can replace a built-in mode tool without replacing provider state.""" + session = AgentSession(session_id="session-1") + mode_provider = AgentModeProvider(source_id="ui_mode", expose_mode_set=False) + + @tool + def update_mode(mode: str) -> str: + """Update the application's mode.""" + return set_agent_mode( + session, + mode, + source_id=mode_provider.source_id, + available_modes=mode_provider.available_modes, + notify=False, + ) + + agent = create_harness_agent( + client=_FakeChatClient(), + max_context_window_tokens=128_000, + max_output_tokens=16_384, + mode_provider=mode_provider, + tools=[update_mode], + agent_instructions="Use update_mode for approved mode changes.", + disable_todo=True, + disable_file_memory=True, + disable_web_search=True, + ) + _, options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Start planning"])], + ) + tools = options["tools"] + assert isinstance(tools, list) + assert [mode_tool.name for mode_tool in tools if isinstance(mode_tool, FunctionTool)] == ["update_mode", "mode_get"] + assert "Use update_mode for approved mode changes." in options["instructions"] + assert "mode_set" not in options["instructions"] + replacement = next( + mode_tool for mode_tool in tools if isinstance(mode_tool, FunctionTool) and mode_tool.name == "update_mode" + ) + await replacement.invoke(arguments={"mode": "execute"}) + assert get_agent_mode(session, source_id=mode_provider.source_id) == "execute" + + updated_context, updated_options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Continue"])], + ) + assert "You are currently operating in the execute mode." in updated_options["instructions"] + assert updated_context.context_messages.get(mode_provider.source_id, []) == [] + + def test_create_harness_agent_disable_file_memory() -> None: """disable_file_memory=True should exclude the FileMemoryProvider.""" agent = create_harness_agent( diff --git a/python/packages/core/tests/core/test_harness_mode.py b/python/packages/core/tests/core/test_harness_mode.py index c4c221f33fe..57120f1ffba 100644 --- a/python/packages/core/tests/core/test_harness_mode.py +++ b/python/packages/core/tests/core/test_harness_mode.py @@ -11,11 +11,13 @@ Agent, AgentModeProvider, AgentSession, + FunctionTool, Message, SupportsChatGetResponse, get_agent_mode, set_agent_mode, ) +from agent_framework._harness._mode import DEFAULT_MODE_INSTRUCTIONS, DEFAULT_MODE_MAP def _tool_by_name(tools: list[object], name: str) -> object: @@ -210,6 +212,132 @@ async def test_agent_mode_context_provider_updates_agent_mode( assert set_agent_mode(session, "plan", source_id=provider.source_id) == "plan" +@pytest.mark.parametrize("expose_mode_set", [True, False]) +@pytest.mark.parametrize("expose_mode_get", [True, False]) +async def test_agent_mode_provider_tool_exposure( + chat_client_base: SupportsChatGetResponse, expose_mode_set: bool, expose_mode_get: bool +) -> None: + """Tool exposure must match built-in guidance without disabling the mode workflow.""" + session = AgentSession(session_id="session-1") + provider = AgentModeProvider(expose_mode_set=expose_mode_set, expose_mode_get=expose_mode_get) + agent = Agent(client=chat_client_base, context_providers=[provider]) + + _, options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Start planning"])], + ) + tools: list[object] = options.get("tools") or [] + expected_names = [ + name for name, exposed in (("mode_set", expose_mode_set), ("mode_get", expose_mode_get)) if exposed + ] + assert [tool.name for tool in tools if isinstance(tool, FunctionTool)] == expected_names + instructions = options["instructions"] + assert isinstance(instructions, str) + assert ("mode_set" in instructions) == expose_mode_set + assert ("mode_get" in instructions) == expose_mode_get + assert "### Mandatory Mode based Workflow" in instructions + assert "get user approval before proceeding" in instructions + assert "You are currently operating in the plan mode." in instructions + assert get_agent_mode(session) == "plan" + if not expose_mode_set: + assert "only after the mode has changed" in instructions + for mode_tool in tools: + assert isinstance(mode_tool, FunctionTool) + assert mode_tool.approval_mode == "never_require" + if mode_tool.name == "mode_set": + result = await mode_tool.invoke(arguments={"mode": "execute"}) + assert result[0].text is not None + assert json.loads(result[0].text) == {"mode": "execute", "message": "Mode changed to 'execute'."} + else: + result = await mode_tool.invoke() + assert result[0].text is not None + assert json.loads(result[0].text) == {"mode": "execute" if expose_mode_set else "plan"} + + +@pytest.mark.parametrize("instructions", [None, ""]) +def test_agent_mode_provider_preserves_default_instructions(instructions: str | None) -> None: + """Suppression on one provider must not alter another provider's defaults.""" + AgentModeProvider(expose_mode_set=False, expose_mode_get=False) + provider = AgentModeProvider(instructions=instructions) + mode_lines = "".join(f"#### {name}\n\n{text}\n\n" for name, text in DEFAULT_MODE_MAP.items()) + expected = DEFAULT_MODE_INSTRUCTIONS.replace("{available_modes}", mode_lines).replace("{current_mode}", "plan") + assert provider._build_instructions("plan") == expected # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("custom_instructions", [True, False]) +@pytest.mark.parametrize("custom_modes", [True, False]) +def test_agent_mode_provider_preserves_custom_instructions(custom_instructions: bool, custom_modes: bool) -> None: + """Only built-in guidance should be adapted; caller text still expands placeholders.""" + mode_text = "Custom mode_get and mode_set guidance for {current_mode}." + mode_map = {"draft": mode_text} if custom_modes else None + instructions = "Use update_mode, not mode_set or mode_get. {current_mode}\n{available_modes}" + provider = AgentModeProvider( + expose_mode_set=False, + expose_mode_get=False, + instructions=instructions if custom_instructions else None, + mode_instructions=mode_map, + ) + current_mode = "draft" if custom_modes else "plan" + rendered = provider._build_instructions(current_mode) # pyright: ignore[reportPrivateUsage] + if custom_modes: + assert mode_map == {"draft": mode_text} + assert f"Custom mode_get and mode_set guidance for {current_mode}." in rendered + if custom_instructions: + assert rendered.startswith(f"Use update_mode, not mode_set or mode_get. {current_mode}\n") + assert "{available_modes}" not in rendered + assert "{current_mode}" not in rendered + if not custom_modes and not custom_instructions: + assert "mode_set" not in rendered + assert "mode_get" not in rendered + + +async def test_agent_mode_provider_hidden_tools_preserve_state_and_notifications( + chat_client_base: SupportsChatGetResponse, +) -> None: + """State and one-shot external notifications must survive even with no mode tools.""" + session = AgentSession(session_id="session-1") + provider = AgentModeProvider( + source_id="ui_mode", default_mode="execute", expose_mode_set=False, expose_mode_get=False + ) + agent = Agent(client=chat_client_base, context_providers=[provider]) + _, first_options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Start"])], + ) + assert not first_options.get("tools") + assert "You are currently operating in the execute mode." in first_options["instructions"] + assert get_agent_mode(session, source_id=provider.source_id) == "execute" + set_agent_mode(session, "plan", source_id=provider.source_id) + + changed_context, changed_options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Continue"])], + ) + assert "You are currently operating in the plan mode." in changed_options["instructions"] + notifications = changed_context.context_messages.get(provider.source_id, []) + assert len(notifications) == 1 + assert notifications[0].role == "user" + assert 'from "execute" to "plan"' in notifications[0].text + assert "previous_mode_for_notification" not in session.state[provider.source_id] + + set_agent_mode(session, "plan", source_id=provider.source_id) + unchanged_context, _ = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Continue planning"])], + ) + assert unchanged_context.context_messages.get(provider.source_id, []) == [] + + reconfigured_provider = AgentModeProvider(source_id=provider.source_id, default_mode="execute") + reconfigured_agent = Agent(client=chat_client_base, context_providers=[reconfigured_provider]) + _, restored_options = await reconfigured_agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Status"])], + ) + assert "You are currently operating in the plan mode." in restored_options["instructions"] + assert get_agent_mode(session, source_id=provider.source_id) == "plan" + assert DEFAULT_MODE_SOURCE_ID not in session.state + + def test_default_mode_falls_back_to_first_available_mode() -> None: """When ``default_mode`` is omitted, helpers and provider should use the first configured mode.""" session = AgentSession(session_id="session-1") @@ -251,6 +379,29 @@ def test_set_agent_mode_no_op_does_not_record_previous_mode() -> None: assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID] +def test_set_agent_mode_can_skip_external_change_notification() -> None: + """Agent-invoked replacement tools should be able to avoid a redundant notification.""" + session = AgentSession(session_id="session-1") + set_agent_mode(session, "plan") + set_agent_mode(session, "execute", notify=False) + + assert get_agent_mode(session) == "execute" + assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID] + + +def test_set_agent_mode_without_notification_clears_pending_notification() -> None: + """An agent-observed update should replace pending external transition context.""" + session = AgentSession(session_id="session-1") + set_agent_mode(session, "plan") + set_agent_mode(session, "execute") + assert session.state[DEFAULT_MODE_SOURCE_ID]["previous_mode_for_notification"] == "plan" + + set_agent_mode(session, "plan", notify=False) + + assert get_agent_mode(session) == "plan" + assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID] + + async def test_agent_mode_provider_injects_user_message_after_external_change( chat_client_base: SupportsChatGetResponse, ) -> None: diff --git a/python/samples/02-agents/context_providers/README.md b/python/samples/02-agents/context_providers/README.md index df0a3a3b1cc..1c8b5990146 100644 --- a/python/samples/02-agents/context_providers/README.md +++ b/python/samples/02-agents/context_providers/README.md @@ -57,3 +57,30 @@ These samples demonstrate how to use context providers to enrich agent conversat - Azure CLI authentication (`az login`) See each subfolder's README for provider-specific prerequisites. + +## Application-controlled modes + +Configure `AgentModeProvider(expose_mode_set=False)` to hide the built-in setter +when your application owns mode changes. Use `expose_mode_get=False` to hide the +getter independently, or set both flags to `False` to expose neither tool. Both +flags default to `True`. Mode state, per-turn workflow instructions, and external +mode-change notifications remain active even when both tools are hidden. + +Pass the configured provider through `create_harness_agent(mode_provider=...)` +(or `Agent(context_providers=[...])`). Supply any replacement tool, such as +`update_mode`, through the existing `tools` argument. Keep that tool and your UI +on the same session-backed state by using `get_agent_mode` and `set_agent_mode` +with the provider's `source_id` and `available_modes`, and its `default_mode` when +reading. A replacement tool should call `set_agent_mode(..., notify=False)` +because the agent already observes the tool result; this also clears any pending +external-change notification. UI-driven changes retain the default `notify=True` +so the agent sees the external change on its next run. Do not disable the entire +provider with `disable_mode=True`. + +Built-in guidance only advertises enabled tools. Without the built-in setter, +it defers approved transitions to the application's configured mode-change +mechanism. Supply replacement-specific guidance through the existing +`instructions` and `mode_instructions` options; caller-supplied text is not +rewritten when a tool is hidden, and existing placeholders still expand. +These flags only omit tools contributed by this provider: they do not filter +application-supplied tools or prevent application code from changing mode state. From 63aacbb5c2157c6d6ecc5fe335c4dc7bd9d84b90 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Thu, 17 Sep 2026 11:39:06 +0000 Subject: [PATCH 07/26] Python: expose instrumentation message event control (#8451) --- .../core/agent_framework/observability.py | 11 ++ .../core/tests/core/test_observability.py | 131 ++++++++++++++++++ .../samples/02-agents/observability/README.md | 12 ++ 3 files changed, 154 insertions(+) diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 9ecc90d06c8..fc911223ef9 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -1556,6 +1556,7 @@ def disable_instrumentation() -> None: def enable_instrumentation( *, enable_sensitive_data: bool | None = None, + enable_message_events: bool | None = None, force: bool = False, ) -> None: """Enable instrumentation for Microsoft Agent Framework. @@ -1568,9 +1569,17 @@ def enable_instrumentation( Keyword Args: enable_sensitive_data: Enable OpenTelemetry sensitive events. Overrides the environment variable ENABLE_SENSITIVE_DATA if set. Default is None. + enable_message_events: Emit the baseline v1.36.0 GenAI message events for + model invocation when sensitive data capture is enabled. Explicit values + override the current setting, including values from ENABLE_MESSAGE_EVENTS. + Default is None, which preserves the current setting without re-reading + the environment. Does not affect experimental message span attributes. force: When True, clears any sticky disable previously set by ``disable_instrumentation()`` before enabling. Without it, calls are no-ops if instrumentation has been explicitly disabled. + + Note: + This function does not configure or replace OpenTelemetry providers or exporters. """ global OBSERVABILITY_SETTINGS if OBSERVABILITY_SETTINGS._user_disabled and not force: # type: ignore[reportPrivateUsage] @@ -1587,6 +1596,8 @@ def enable_instrumentation( else: # Re-read from current environment in case env vars were set after import (e.g. load_dotenv()) OBSERVABILITY_SETTINGS.enable_sensitive_data = _read_bool_env("ENABLE_SENSITIVE_DATA") + if enable_message_events is not None: + OBSERVABILITY_SETTINGS.enable_message_events = enable_message_events def configure_otel_providers( diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index 0dcafce37f8..fb4a5b1d38c 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -1914,6 +1914,137 @@ def test_enable_instrumentation_explicit_param_overrides_env(monkeypatch): assert observability.OBSERVABILITY_SETTINGS.enable_sensitive_data is False +@pytest.mark.parametrize("message_events", [False, True]) +def test_enable_instrumentation_message_events_overrides_env(monkeypatch, message_events: bool) -> None: + import agent_framework.observability as observability + + monkeypatch.setenv("ENABLE_MESSAGE_EVENTS", str(not message_events)) + monkeypatch.setenv("ENABLE_SENSITIVE_DATA", "false") + settings = observability.ObservabilitySettings(enable_instrumentation=False) + monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", settings) + assert settings.enable_message_events is not message_events + + observability.enable_instrumentation(enable_message_events=message_events) + + assert settings.enable_instrumentation is True + assert settings.enable_message_events is message_events + assert settings.enable_sensitive_data is False + + +@pytest.mark.parametrize("message_events", [False, True]) +@pytest.mark.parametrize("source", ["environment", "providers", "instrumentation"]) +def test_enable_instrumentation_preserves_message_events(monkeypatch, message_events: bool, source: str) -> None: + import agent_framework.observability as observability + + monkeypatch.setenv("ENABLE_MESSAGE_EVENTS", str(message_events if source == "environment" else not message_events)) + settings = observability.ObservabilitySettings() + monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", settings) + if source == "providers": + with patch.object(settings, "_configure"): + observability.configure_otel_providers(enable_message_events=message_events) + elif source == "instrumentation": + observability.enable_instrumentation(enable_message_events=message_events) + + for env_value in [str(not message_events), None]: + if env_value is None: + monkeypatch.delenv("ENABLE_MESSAGE_EVENTS") + else: + monkeypatch.setenv("ENABLE_MESSAGE_EVENTS", env_value) + observability.enable_instrumentation() + assert settings.enable_message_events is message_events + observability.enable_instrumentation(enable_message_events=None) + assert settings.enable_message_events is message_events + + +@pytest.mark.parametrize("current_value", [False, True]) +@pytest.mark.parametrize("message_events", [False, True, None]) +@pytest.mark.parametrize("force", [False, True]) +def test_enable_instrumentation_message_events_respects_sticky_disable( + monkeypatch, current_value: bool, message_events: bool | None, force: bool +) -> None: + import agent_framework.observability as observability + + settings = observability.ObservabilitySettings(enable_message_events=current_value) + monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", settings) + observability.disable_instrumentation() + + observability.enable_instrumentation(enable_message_events=message_events, enable_sensitive_data=True, force=force) + + assert settings.enable_instrumentation is force + assert settings.enable_sensitive_data is force + assert settings.is_user_disabled is not force + expected = message_events if force and message_events is not None else current_value + assert settings.enable_message_events is expected + + +@pytest.mark.parametrize("is_setup", [False, True]) +def test_enable_instrumentation_message_events_preserves_providers( + monkeypatch, span_exporter, log_record_exporter, is_setup: bool +) -> None: + from opentelemetry import metrics, trace + from opentelemetry._logs import get_logger_provider + + import agent_framework.observability as observability + + settings = observability.OBSERVABILITY_SETTINGS + monkeypatch.setattr(settings, "enable_console_exporters", True) + monkeypatch.setattr(settings, "_executed_setup", is_setup) + providers = (trace.get_tracer_provider(), get_logger_provider(), metrics.get_meter_provider()) + with patch.object(settings, "_configure") as configure: + for message_events in [False, True, None]: + observability.enable_instrumentation(enable_message_events=message_events) + assert settings.is_setup is is_setup + assert settings.enable_console_exporters is True + assert trace.get_tracer_provider() is providers[0] + assert get_logger_provider() is providers[1] + assert metrics.get_meter_provider() is providers[2] + configure.assert_not_called() + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("message_events", [False, True]) +@pytest.mark.parametrize("enable_sensitive_data", [False, True], indirect=True) +@pytest.mark.parametrize("semconv", ["", "gen_ai_latest_experimental"]) +async def test_enable_instrumentation_controls_message_event_emission( + mock_chat_client, + span_exporter: InMemorySpanExporter, + log_record_exporter, + enable_sensitive_data: bool, + stream: bool, + message_events: bool, + semconv: str, +) -> None: + import agent_framework.observability as observability + + observability.OBSERVABILITY_SETTINGS.otel_semconv_stability_opt_in = semconv + observability.enable_instrumentation(enable_message_events=message_events) + client = mock_chat_client() + messages = [Message("user", ["Test message"])] + if stream: + response_stream = client.get_response(messages=messages, options={"model": "Test"}, stream=True) + async for _ in response_stream: + pass + await response_stream.get_final_response() + else: + await client.get_response(messages=messages, options={"model": "Test"}) + + records = [record.log_record for record in log_record_exporter.get_finished_logs()] + expected_events = [] + if message_events and enable_sensitive_data: + expected_events = [OtelAttr.USER_MESSAGE.value] + if stream: + expected_events.append(OtelAttr.CHOICE.value) + assert [record.event_name for record in records] == expected_events + if records: + assert records[0].body == {"content": "Test message"} + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + attributes = spans[0].attributes or {} + assert (OtelAttr.INPUT_MESSAGES in attributes) is (enable_sensitive_data and bool(semconv)) + assert (OtelAttr.OUTPUT_MESSAGES in attributes) is (enable_sensitive_data and bool(semconv)) + + def test_enable_instrumentation_does_not_touch_console_exporters(monkeypatch): """Test enable_instrumentation does not modify enable_console_exporters (it is an exporter concern).""" import importlib diff --git a/python/samples/02-agents/observability/README.md b/python/samples/02-agents/observability/README.md index 822006b17dd..e69af64f684 100644 --- a/python/samples/02-agents/observability/README.md +++ b/python/samples/02-agents/observability/README.md @@ -88,6 +88,18 @@ from microsoft.opentelemetry import use_microsoft_opentelemetry use_microsoft_opentelemetry(enable_azure_monitor=True) ``` +To disable Agent Framework's baseline GenAI message events without changing providers or exporters configured by a third party, use the instrumentation-only API: + +```python +from agent_framework.observability import enable_instrumentation + +enable_instrumentation(enable_message_events=False) +``` + +An explicit `True` or `False` overrides the current message-event setting, including a value read from `ENABLE_MESSAGE_EVENTS`. Omitting the argument or passing `None` preserves the current setting without re-reading the environment, including settings established by `configure_otel_providers()` or an earlier explicit call. + +This flag controls baseline v1.36.0 events such as `gen_ai.user.message` and `gen_ai.choice`; it does not disable experimental message span attributes, spans, or metrics. Message events still require sensitive-data capture to be enabled separately. Existing `enable_sensitive_data` behavior is unchanged, and [sticky disable](#disabling-instrumentation) still requires `force=True` to re-enable instrumentation. + ```python from azure.monitor.opentelemetry import configure_azure_monitor from agent_framework.observability import create_resource, enable_sensitive_telemetry From 0275ebfd026e3549b9377788aedb58f424997fb2 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Thu, 17 Sep 2026 11:46:37 +0000 Subject: [PATCH 08/26] .NET: [BREAKING] Scope provider-backed MCP sessions per invocation (#8425) * Scope provider-backed MCP sessions to each invocation Create and clean up separate MCP sessions for custom client-provider calls in .NET and Python. Preserve default session caching and caller-owned HTTP clients, and add lifecycle coverage and compatibility guidance. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Make cleanup cancellation test compatible with Python 3.10 Assert cancellation and resource cleanup without depending on a cancellation message that older asyncio task implementations may consume. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Handle reentrant MCP shutdown and cleanup failures Reject self-draining shutdown across inherited invocation contexts. Preserve invocation outcomes when session or transport cleanup fails, report non-cancellation cleanup failures separately, and add lifecycle regression coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../DefaultMcpToolHandler.cs | 237 +++++- .../README.md | 27 +- .../DefaultMcpToolHandlerLifetimeTests.cs | 694 ++++++++++++++++ python/packages/declarative/AGENTS.md | 15 + python/packages/declarative/README.md | 22 + .../_workflows/_mcp_handler.py | 193 +++-- .../tests/test_default_mcp_tool_handler.py | 745 ++++++++++++++++-- 7 files changed, 1772 insertions(+), 161 deletions(-) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs index 49541ebf428..b9c465f9f03 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative.Mcp/DefaultMcpToolHandler.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics; using System.IO; using System.Linq; using System.Net.Http; @@ -25,6 +26,11 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp; /// This provider supports per-server authentication via the httpClientProvider callback. /// The callback allows different MCP servers to use different authentication configurations by returning /// a pre-configured for each server. +/// Provider-backed invocations create and dispose a separate MCP session for every call, including +/// tools/list, because provider authentication is not represented in the session cache key. +/// Without a provider, sessions are cached by server URL, label, connection name, and explicit headers. +/// Non-cancellation cleanup failures are reported through warnings without replacing +/// the invocation result or error. /// public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable { @@ -39,9 +45,14 @@ public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable private static readonly JsonWriterOptions s_toolListJsonWriterOptions = new() { Indented = true }; private readonly Func>? _httpClientProvider; - private readonly Dictionary<(string Url, string Label, string Connection, string HeadersHash), McpClient> _clients = []; + private readonly Func _httpMessageHandlerFactory; + private readonly Dictionary<(string Url, string Label, string Connection, string HeadersHash), ClientConnection> _clients = []; private readonly Dictionary _ownedHttpClients = []; private readonly SemaphoreSlim _clientLock = new(1, 1); + private readonly AsyncLocal _providerInvocationContext = new(); + private TaskCompletionSource? _providerInvocationsDrained; + private int _activeProviderInvocations; + private bool _disposing; /// /// Initializes a new instance of the class. @@ -51,6 +62,13 @@ public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable /// The callback receives (serverUrl, cancellationToken) and should return an HttpClient /// configured with any required authentication. Return to use a default HttpClient with no auth. /// + /// The callback is invoked for each tool invocation. MCP sessions are not cached when a callback is + /// configured, even if it returns the same client or . This adds session setup + /// per call and does not preserve server-side session state between calls. Supplied HTTP clients + /// remain caller-owned. Applications requiring session continuity should provide an + /// with an explicitly scoped authentication and session lifetime. + /// + /// /// Security: HttpClients created by this handler pin credential headers to the configured server /// origin and disable auto-redirect. When you supply your own , you are /// responsible for equivalent protection — attach credentials only for the configured server origin @@ -59,8 +77,16 @@ public sealed class DefaultMcpToolHandler : IMcpToolHandler, IAsyncDisposable /// /// public DefaultMcpToolHandler(Func>? httpClientProvider = null) + : this(httpClientProvider, CreateHttpMessageHandler) + { + } + + internal DefaultMcpToolHandler( + Func>? httpClientProvider, + Func httpMessageHandlerFactory) { this._httpClientProvider = httpClientProvider; + this._httpMessageHandlerFactory = Throw.IfNull(httpMessageHandlerFactory); } /// @@ -76,12 +102,70 @@ public async Task InvokeToolAsync( if (IsListToolsToolName(toolName)) { ThrowIfListToolsArgumentsSpecified(arguments); - McpClient listToolsClient = await this.GetOrCreateClientAsync(serverUrl, serverLabel, headers, connectionName, cancellationToken).ConfigureAwait(false); - IList tools = await listToolsClient.ListToolsAsync(cancellationToken: cancellationToken).ConfigureAwait(false); - return CreateListToolsResultContent(tools.Select(tool => tool.ProtocolTool)); + } + + if (this._httpClientProvider is not null) + { + TaskCompletionSource invocationCompleted = new(TaskCreationOptions.RunContinuationsAsynchronously); + await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + this.ThrowIfDisposing(); + if (this._activeProviderInvocations++ == 0) + { + this._providerInvocationsDrained = new(TaskCreationOptions.RunContinuationsAsynchronously); + } + } + finally + { + this._clientLock.Release(); + } + + ProviderInvocationContext? previousInvocation = this._providerInvocationContext.Value; + this._providerInvocationContext.Value = new(invocationCompleted.Task, previousInvocation); + try + { + ClientConnection invocationClient = await this.CreateClientAsync( + serverUrl.Trim(), serverLabel, headers, httpClientCacheKey: null, cancellationToken).ConfigureAwait(false); + await using (invocationClient.ConfigureAwait(false)) + { + return await InvokeClientAsync(invocationClient.Client, toolName, arguments, cancellationToken).ConfigureAwait(false); + } + } + finally + { + await this._clientLock.WaitAsync(CancellationToken.None).ConfigureAwait(false); + try + { + if (--this._activeProviderInvocations == 0) + { + this._providerInvocationsDrained!.SetResult(true); + } + } + finally + { + invocationCompleted.SetResult(true); + this._providerInvocationContext.Value = previousInvocation; + this._clientLock.Release(); + } + } } McpClient client = await this.GetOrCreateClientAsync(serverUrl, serverLabel, headers, connectionName, cancellationToken).ConfigureAwait(false); + return await InvokeClientAsync(client, toolName, arguments, cancellationToken).ConfigureAwait(false); + } + + private static async Task InvokeClientAsync( + McpClient client, + string toolName, + IDictionary? arguments, + CancellationToken cancellationToken) + { + if (IsListToolsToolName(toolName)) + { + IList tools = await client.ListToolsAsync(cancellationToken: cancellationToken).ConfigureAwait(false); + return CreateListToolsResultContent(tools.Select(tool => tool.ProtocolTool)); + } McpServerToolResultContent resultContent = new(Guid.NewGuid().ToString()); @@ -119,12 +203,42 @@ internal static McpServerToolResultContent CreateListToolsResultContent(IEnumera } /// + /// + /// Disposal is requested from an active provider-backed invocation, including a child execution context. + /// Dispose the handler from its owning scope instead. + /// public async ValueTask DisposeAsync() { + for (ProviderInvocationContext? context = this._providerInvocationContext.Value; context is not null; context = context.Parent) + { + if (!context.Completion.IsCompleted) + { + throw new InvalidOperationException("Cannot dispose the MCP handler from an active provider-backed invocation."); + } + } + + Task? providerInvocations; await this._clientLock.WaitAsync().ConfigureAwait(false); try { - foreach (McpClient client in this._clients.Values) + this.ThrowIfDisposing(); + this._disposing = true; + providerInvocations = this._providerInvocationsDrained?.Task; + } + finally + { + this._clientLock.Release(); + } + + if (providerInvocations is not null) + { + await providerInvocations.ConfigureAwait(false); + } + + await this._clientLock.WaitAsync().ConfigureAwait(false); + try + { + foreach (ClientConnection client in this._clients.Values) { await client.DisposeAsync().ConfigureAwait(false); } @@ -147,6 +261,16 @@ public async ValueTask DisposeAsync() this._clientLock.Dispose(); } + [System.Diagnostics.CodeAnalysis.SuppressMessage("Maintainability", "CA1513:Use ObjectDisposedException throw helper", + Justification = "The helper is not available on .NET Framework or .NET Standard 2.0.")] + private void ThrowIfDisposing() + { + if (this._disposing) + { + throw new ObjectDisposedException(nameof(DefaultMcpToolHandler)); + } + } + private async Task GetOrCreateClientAsync( string serverUrl, string? serverLabel, @@ -160,14 +284,15 @@ private async Task GetOrCreateClientAsync( await this._clientLock.WaitAsync(cancellationToken).ConfigureAwait(false); try { - if (this._clients.TryGetValue(clientCacheKey, out McpClient? existingClient)) + this.ThrowIfDisposing(); + if (this._clients.TryGetValue(clientCacheKey, out ClientConnection? existingClient)) { - return existingClient; + return existingClient.Client; } - McpClient newClient = await this.CreateClientAsync(trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); + ClientConnection newClient = await this.CreateClientAsync(trimmedUrl, serverLabel, headers, trimmedUrl, cancellationToken).ConfigureAwait(false); this._clients[clientCacheKey] = newClient; - return newClient; + return newClient.Client; } finally { @@ -188,36 +313,25 @@ internal static (string Url, string Label, string Connection, string HeadersHash IDictionary? headers) => (trimmedUrl, serverLabel ?? string.Empty, connectionName ?? string.Empty, ComputeHeadersHash(headers)); - private async Task CreateClientAsync( + private async Task CreateClientAsync( string serverUrl, string? serverLabel, IDictionary? headers, - string httpClientCacheKey, + string? httpClientCacheKey, CancellationToken cancellationToken) { - // Get or create HttpClient (Can be shared across McpClients for the same server) + // Only the no-provider path shares handler-owned HTTP clients. HttpClient? httpClient = null; + bool ownsHttpClient = false; if (this._httpClientProvider is not null) { httpClient = await this._httpClientProvider(serverUrl, cancellationToken).ConfigureAwait(false); } - if (httpClient is null && !this._ownedHttpClients.TryGetValue(httpClientCacheKey, out httpClient)) + if (httpClient is null && + (httpClientCacheKey is null || !this._ownedHttpClients.TryGetValue(httpClientCacheKey, out httpClient))) { - // Disable cookies so handler-level state (cookie jar) cannot cross the cache-key - // isolation boundary established by GetOrCreateClientAsync. The actual MCP auth - // travels via AdditionalHeaders (set per-transport below), not session cookies. - // Disable auto-redirect so an HTTP redirect cannot carry credential headers to a - // different origin (defense-in-depth alongside the origin pinning applied below). - // CheckCertificateRevocationList satisfies CA5399 since we're explicitly constructing the handler. - HttpClientHandler handler = new() - { - UseCookies = false, - AllowAutoRedirect = false, - CheckCertificateRevocationList = true - }; - // Pin credential headers to the configured server origin as defense-in-depth. Forcing // StreamableHttp (below) already removes the primary vector (a server-advertised cross-origin // SSE message endpoint), and AllowAutoRedirect=false blocks auto-redirects. This handler is the @@ -225,9 +339,16 @@ private async Task CreateClientAsync( // origin even if a future change re-enables AutoDetect or redirects, or the SDK constructs a // request to a new URI (AdditionalHeaders are re-stamped by the transport, so HttpClient's own // redirect header-stripping does not cover them). - OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = handler }; + OriginPinningHandler pinningHandler = new(new Uri(serverUrl)) { InnerHandler = this._httpMessageHandlerFactory() }; httpClient = new HttpClient(pinningHandler); - this._ownedHttpClients[httpClientCacheKey] = httpClient; + if (httpClientCacheKey is null) + { + ownsHttpClient = true; + } + else + { + this._ownedHttpClients[httpClientCacheKey] = httpClient; + } } HttpClientTransportOptions transportOptions = new() @@ -242,9 +363,65 @@ private async Task CreateClientAsync( TransportMode = HttpTransportMode.StreamableHttp }; - HttpClientTransport transport = new(transportOptions, httpClient); + HttpClientTransport transport = new(transportOptions, httpClient, ownsHttpClient: ownsHttpClient); + + try + { + McpClient client = await McpClient.CreateAsync(transport, cancellationToken: cancellationToken).ConfigureAwait(false); + return new ClientConnection(client, transport); + } + catch + { + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + throw; + } + } + + private static HttpMessageHandler CreateHttpMessageHandler() => + new HttpClientHandler + { + // Keep cookies out of pooled transport state and prevent credential-bearing redirects. + UseCookies = false, + AllowAutoRedirect = false, + CheckCertificateRevocationList = true + }; + + private sealed class ProviderInvocationContext(Task completion, ProviderInvocationContext? parent) + { + public Task Completion { get; } = completion; + + public ProviderInvocationContext? Parent { get; } = parent; + } + + internal sealed class ClientConnection(McpClient client, IAsyncDisposable transport) : IAsyncDisposable + { + public McpClient Client { get; } = client; + + public async ValueTask DisposeAsync() + { + try + { + await DisposeResourceAsync(this.Client, "session").ConfigureAwait(false); + } + finally + { + // McpClient owns the connected session, not the reusable transport factory. + await DisposeResourceAsync(transport, "transport").ConfigureAwait(false); + } + } + } - return await McpClient.CreateAsync(transport, cancellationToken: cancellationToken).ConfigureAwait(false); + private static async ValueTask DisposeResourceAsync(T resource, string resourceName) + where T : IAsyncDisposable + { + try + { + await resource.DisposeAsync().ConfigureAwait(false); + } + catch (Exception exception) when (exception is not OperationCanceledException) + { + Trace.TraceWarning("Failed to dispose MCP {0}: {1}", resourceName, exception); + } } /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/README.md b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/README.md index 2a50b6045d2..fcfb73ee1a2 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/README.md +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/README.md @@ -13,6 +13,31 @@ we've provided a console application that is able to execute any declarative wor ## Actions +### MCP session lifetime + +`DefaultMcpToolHandler` caches MCP sessions by server URL, label, connection name, +and explicit request headers when no custom HTTP client provider is configured. +Dispose the handler when its owning scope ends. + +When `httpClientProvider` is configured, each invocation (including `tools/list`) +calls the provider, opens a separate MCP session, and disposes that session before +completing. This also applies when the provider returns the same `HttpClient` or +returns `null` for the default transport. Caller-supplied HTTP clients remain +caller-owned; handler-created fallback clients are disposed with the invocation. +Handler disposal waits for active provider-backed invocations to finish cleanup. +Disposal from within an active invocation (including a provider callback or its +child tasks) throws `InvalidOperationException` without starting shutdown, rather +than waiting on itself. Dispose the handler from its owning scope. +Non-cancellation session and transport cleanup failures are reported through +`System.Diagnostics.Trace` warnings and do not replace a tool result or its +original error. + +Provider-backed calls therefore incur session setup per invocation and do not +preserve server-side session state between calls. Applications requiring session +continuity should implement `IMcpToolHandler` with an explicit authentication +ownership and session lifetime contract, rather than sharing sessions based only +on the server address. + ### ⚙ī¸ Foundry Actions |Action|Description| @@ -56,5 +81,3 @@ we've provided a console application that is able to execute any declarative wor |**EndWorkflow**|Ends the current workflow or sub-workflow within a broader conversation flow. This helps modularize complex interactions. |**Foreach**|Iterates through a collection of items, executing a set of actions for each. Ideal for processing lists or batch operations. |**GotoAction**|Jumps directly to a specified action within the workflow. Enables non-linear navigation in the logic flow. - - diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs new file mode 100644 index 00000000000..3a481fcd0e6 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests/DefaultMcpToolHandlerLifetimeTests.cs @@ -0,0 +1,694 @@ +īģŋ// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Globalization; +using System.Net; +using System.Net.Http; +using System.Text; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using ModelContextProtocol.Client; +using Moq; +using Moq.Protected; + +namespace Microsoft.Agents.AI.Workflows.Declarative.Mcp.UnitTests; + +/// +/// Functional session-lifetime tests using an in-memory HTTP transport and inert tool results. +/// +[System.Diagnostics.CodeAnalysis.SuppressMessage("Reliability", "CA2025:Do not pass disposable objects into unawaited tasks", + Justification = "Provider callbacks return completed tasks; each invocation and handler disposal is awaited before disposing caller-owned clients.")] +public sealed class DefaultMcpToolHandlerLifetimeTests +{ + [Theory] + [InlineData("ping")] + [InlineData(DefaultMcpToolHandler.ListToolsToolName)] + public async Task Provider_SequentialInvocations_CreateAndDisposeSeparateSessionsAsync(string toolName) + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + int providerCalls = 0; + DefaultMcpToolHandler handler = new((_, _) => + { + providerCalls++; + return Task.FromResult(client); + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeAsync(handler, toolName, timeout.Token); + Assert.Equal(1, stub.Terminations); + await InvokeAsync(handler, toolName, timeout.Token); + await handler.DisposeAsync(); + + // Assert + Assert.Equal(2, providerCalls); + Assert.Equal(2, stub.Initializations); + Assert.Equal(2, stub.Terminations); + stub.Handlers[0].Protected().Verify("Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task ProviderReturningNull_DisposesFallbackTransportPerInvocationAsync() + { + // Arrange + ProtocolStub stub = new(); + int providerCalls = 0; + await using DefaultMcpToolHandler handler = new((_, _) => + { + providerCalls++; + return Task.FromResult(null); + }, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeAsync(handler, "ping", timeout.Token); + await InvokeAsync(handler, "ping", timeout.Token); + + // Assert + Assert.Equal(2, providerCalls); + Assert.Equal(2, stub.Initializations); + Assert.Equal(2, stub.Terminations); + Assert.Equal(2, stub.Handlers.Count); + foreach (Mock transport in stub.Handlers) + { + transport.Protected().Verify("Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + } + + [Fact] + public async Task NoProvider_ReusesSessionUntilHandlerDisposalAsync() + { + // Arrange + ProtocolStub stub = new(); + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeAsync(handler, "ping", timeout.Token); + await InvokeAsync(handler, DefaultMcpToolHandler.ListToolsToolName, timeout.Token); + + // Assert + Assert.Equal(1, stub.Initializations); + Assert.Equal(0, stub.Terminations); + Mock transport = Assert.Single(stub.Handlers); + transport.Protected().Verify("Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); + + await handler.DisposeAsync(); + Assert.Equal(1, stub.Terminations); + transport.Protected().Verify("Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task NoProvider_DifferentConnectionNames_UseSeparateCachedSessionsAsync() + { + // Arrange + ProtocolStub stub = new(); + DefaultMcpToolHandler handler = new(null, stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + foreach (string connectionName in new[] { "first", "second", "first" }) + { + await handler.InvokeToolAsync( + "https://mcp.example/api", null, "ping", null, null, connectionName, timeout.Token); + } + + // Assert + Assert.Equal(2, stub.Initializations); + Assert.Single(stub.Handlers); + Assert.Equal(0, stub.Terminations); + await handler.DisposeAsync(); + Assert.Equal(2, stub.Terminations); + } + + [Fact] + public async Task Provider_Failure_DoesNotPreventHandlerDisposalAsync() + { + // Arrange + int providerCalls = 0; + DefaultMcpToolHandler handler = new((_, _) => + { + providerCalls++; + throw new InvalidOperationException("Provider unavailable."); + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + InvalidOperationException exception = await Assert.ThrowsAsync( + () => InvokeAsync(handler, "ping", timeout.Token)); + await handler.DisposeAsync(); + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + + // Assert + Assert.Equal("Provider unavailable.", exception.Message); + Assert.Equal(1, providerCalls); + } + + [Fact] + public async Task Provider_ConcurrentInvocations_DoNotCoalesceOrSerializeSessionsAsync() + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + using SemaphoreSlim started = new(0); + using SemaphoreSlim finish = new(0); + stub.BeforeOperationAsync = async token => + { + started.Release(); + await finish.WaitAsync(token); + }; + int providerCalls = 0; + await using DefaultMcpToolHandler handler = new((_, _) => + { + Interlocked.Increment(ref providerCalls); + return Task.FromResult(client); + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task first = InvokeAsync(handler, "ping", timeout.Token); + Task second = InvokeAsync(handler, "ping", timeout.Token); + await started.WaitAsync(timeout.Token); + await started.WaitAsync(timeout.Token); + finish.Release(2); + await Task.WhenAll(first, second); + + // Assert + Assert.Equal(2, providerCalls); + Assert.Equal(2, stub.Initializations); + Assert.Equal(2, stub.Terminations); + } + + [Fact] + public async Task Provider_OperationFailure_DisposesSessionAndPreservesCallerClientAsync() + { + // Arrange + ProtocolStub stub = new() { FailOperation = true }; + using HttpClient client = new(stub.CreateMessageHandler()); + await using DefaultMcpToolHandler handler = new((_, _) => Task.FromResult(client)); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + + // Assert + Assert.Equal(1, stub.Initializations); + Assert.Equal(1, stub.Terminations); + stub.Handlers[0].Protected().Verify("Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Provider_InitializationFailure_RespectsTransportOwnershipAsync(bool provideClient) + { + // Arrange + ProtocolStub stub = new() { FailInitialization = true }; + using HttpClient? client = provideClient ? new(stub.CreateMessageHandler()) : null; + await using DefaultMcpToolHandler handler = new( + (_, _) => Task.FromResult(client), stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + + // Assert + Mock transport = Assert.Single(stub.Handlers); + transport.Protected().Verify("Dispose", provideClient ? Times.Never() : Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task Provider_Cancellation_DisposesSessionAsync() + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + using SemaphoreSlim started = new(0); + stub.BeforeOperationAsync = async token => + { + started.Release(); + await Task.Delay(Timeout.Infinite, token); + }; + await using DefaultMcpToolHandler handler = new((_, _) => Task.FromResult(client)); + using CancellationTokenSource cancellation = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = InvokeAsync(handler, "ping", cancellation.Token); + await started.WaitAsync(cancellation.Token); + cancellation.Cancel(); + await Assert.ThrowsAnyAsync(() => invocation); + + // Assert + Assert.Equal(1, stub.Terminations); + stub.Handlers[0].Protected().Verify("Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task DisposeAsync_WaitsForProviderInvocationCleanupAsync(bool pauseProvider) + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + using SemaphoreSlim started = new(0); + using SemaphoreSlim finish = new(0); + async Task PauseAsync(CancellationToken token) + { + started.Release(); + await finish.WaitAsync(token); + } + + if (!pauseProvider) + { + stub.BeforeOperationAsync = PauseAsync; + } + + int providerCalls = 0; + DefaultMcpToolHandler handler = new(async (_, token) => + { + providerCalls++; + if (pauseProvider) + { + await PauseAsync(token); + } + + return client; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + Task invocation = InvokeAsync(handler, "ping", timeout.Token); + await started.WaitAsync(timeout.Token); + Task disposal = handler.DisposeAsync().AsTask(); + Assert.False(disposal.IsCompleted); + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + finish.Release(); + await invocation; + await disposal; + + // Assert + Assert.Equal(1, providerCalls); + Assert.Equal(1, stub.Terminations); + stub.Handlers[0].Protected().Verify("Dispose", Times.Never(), ItExpr.Is(disposing => disposing)); + } + + [Fact] + public async Task Provider_InvalidDiscoveryArguments_DoesNotCallProviderAsync() + { + // Arrange + int providerCalls = 0; + await using DefaultMcpToolHandler handler = new((_, _) => + { + providerCalls++; + return Task.FromResult(null); + }); + + // Act + await Assert.ThrowsAsync(() => handler.InvokeToolAsync( + "https://mcp.example/api", null, DefaultMcpToolHandler.ListToolsToolName, + new Dictionary { ["unused"] = true }, null, null)); + + // Assert + Assert.Equal(0, providerCalls); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Provider_ReentrantDisposal_ThrowsWithoutStartingShutdownAsync(bool fromChildTask) + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + DefaultMcpToolHandler? handler = null; + int rejectedDisposals = 0; + handler = new(async (_, _) => + { + async Task DisposeAsync() => await handler!.DisposeAsync(); + InvalidOperationException error = await Assert.ThrowsAsync( + () => fromChildTask ? Task.Run(DisposeAsync) : DisposeAsync()); + Assert.Contains("active provider-backed invocation", error.Message); + rejectedDisposals++; + return client; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeAsync(handler, "ping", timeout.Token); + await InvokeAsync(handler, "ping", timeout.Token); + await handler.DisposeAsync(); + + // Assert + Assert.Equal(2, rejectedDisposals); + Assert.Equal(2, stub.Terminations); + } + + [Fact] + public async Task Provider_CompletedInheritedContext_CanDisposeHandlerAsync() + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + using SemaphoreSlim finished = new(0); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + Task? inheritedDisposal = null; + DefaultMcpToolHandler? handler = null; + handler = new((_, token) => + { + inheritedDisposal = Task.Run(async () => + { + await finished.WaitAsync(token); + await handler!.DisposeAsync(); + }, token); + return Task.FromResult(client); + }); + + // Act + await InvokeAsync(handler, "ping", timeout.Token); + finished.Release(); + await Assert.IsAssignableFrom(inheritedDisposal); + + // Assert + Assert.Equal(1, stub.Terminations); + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + } + + [Fact] + public async Task Provider_NestedInvocation_RestoresOuterDisposalGuardAsync() + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + bool nested = false; + DefaultMcpToolHandler? handler = null; + handler = new(async (_, token) => + { + if (!nested) + { + nested = true; + await InvokeAsync(handler!, "ping", token); + await Assert.ThrowsAsync(() => handler!.DisposeAsync().AsTask()); + } + + return client; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeAsync(handler, "ping", timeout.Token); + await handler.DisposeAsync(); + + // Assert + Assert.True(nested); + Assert.Equal(2, stub.Terminations); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task Provider_TransportCleanupFailure_PreservesInvocationOutcomeAsync(bool failOperation) + { + // Arrange + ProtocolStub stub = new() { FailOperation = failOperation, FailTransportDisposal = true }; + await using DefaultMcpToolHandler handler = new( + (_, _) => Task.FromResult(null), stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CleanupTraceListener listener = new(); + Trace.Listeners.Add(listener); + try + { + // Act + if (failOperation) + { + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + } + else + { + McpServerToolResultContent result = await InvokeAsync(handler, "ping", timeout.Token); + Assert.Equal("ok", Assert.IsType(Assert.Single(result.Outputs!)).Text); + } + + // Assert + Assert.Equal(1, stub.Terminations); + Assert.Contains("Failed to dispose MCP transport", listener.Output); + Assert.Contains("transport cleanup failed", listener.Output); + Assert.Single(stub.Handlers).Protected().Verify( + "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + finally + { + Trace.Listeners.Remove(listener); + } + } + + [Fact] + public async Task Provider_CompletedNestedContext_StillGuardsActiveOuterInvocationAsync() + { + // Arrange + ProtocolStub stub = new(); + using HttpClient client = new(stub.CreateMessageHandler()); + using SemaphoreSlim nestedFinished = new(0); + Task? inheritedDisposal = null; + int providerCalls = 0; + DefaultMcpToolHandler? handler = null; + handler = new(async (_, token) => + { + if (++providerCalls == 1) + { + await InvokeAsync(handler!, "ping", token); + nestedFinished.Release(); + await Assert.IsAssignableFrom(inheritedDisposal); + } + else + { + inheritedDisposal = Task.Run(async () => + { + await nestedFinished.WaitAsync(token); + await Assert.ThrowsAsync(() => handler!.DisposeAsync().AsTask()); + }, token); + } + + return client; + }); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + + // Act + await InvokeAsync(handler, "ping", timeout.Token); + await handler.DisposeAsync(); + + // Assert + Assert.Equal(2, providerCalls); + Assert.Equal(2, stub.Terminations); + } + + [Fact] + public async Task Provider_InitializationAndCleanupFailure_PreservesInitializationErrorAsync() + { + // Arrange + ProtocolStub stub = new() { FailInitialization = true, FailTransportDisposal = true }; + await using DefaultMcpToolHandler handler = new( + (_, _) => Task.FromResult(null), stub.CreateMessageHandler); + using CancellationTokenSource timeout = new(TimeSpan.FromSeconds(10)); + using CleanupTraceListener listener = new(); + Trace.Listeners.Add(listener); + try + { + // Act + await Assert.ThrowsAsync(() => InvokeAsync(handler, "ping", timeout.Token)); + + // Assert + Assert.Contains("Failed to dispose MCP transport", listener.Output); + Assert.Single(stub.Handlers).Protected().Verify( + "Dispose", Times.AtLeastOnce(), ItExpr.Is(disposing => disposing)); + } + finally + { + Trace.Listeners.Remove(listener); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ClientConnection_SessionCleanupFailure_StillCleansUpTransportAsync(bool cancelSessionCleanup) + { + // Arrange + Mock client = new(); + Exception sessionError = cancelSessionCleanup + ? new OperationCanceledException("session cleanup cancelled") + : new InvalidOperationException("session cleanup failed"); + client.Setup(c => c.DisposeAsync()).Returns(new ValueTask(Task.FromException(sessionError))); + Mock transport = new(); + transport.Setup(t => t.DisposeAsync()).Returns(new ValueTask(Task.FromException( + new InvalidOperationException("transport cleanup failed")))); + DefaultMcpToolHandler.ClientConnection connection = new(client.Object, transport.Object); + using CleanupTraceListener listener = new(); + Trace.Listeners.Add(listener); + try + { + // Act + if (cancelSessionCleanup) + { + Assert.Same(sessionError, await Assert.ThrowsAsync( + () => connection.DisposeAsync().AsTask())); + } + else + { + await connection.DisposeAsync(); + Assert.Contains("Failed to dispose MCP session", listener.Output); + } + + // Assert + Assert.Contains("Failed to dispose MCP transport", listener.Output); + client.Verify(c => c.DisposeAsync(), Times.Once); + transport.Verify(t => t.DisposeAsync(), Times.Once); + } + finally + { + Trace.Listeners.Remove(listener); + } + } + + private static Task InvokeAsync( + DefaultMcpToolHandler handler, string toolName, CancellationToken cancellationToken) => + handler.InvokeToolAsync("https://mcp.example/api", null, toolName, null, null, null, cancellationToken); + + private sealed class ProtocolStub + { + private int _initializations; + private int _terminations; + + public int Initializations => this._initializations; + public int Terminations => this._terminations; + public List> Handlers { get; } = []; + public Func? BeforeOperationAsync { get; set; } + public bool FailInitialization { get; set; } + public bool FailOperation { get; set; } + public bool FailTransportDisposal { get; set; } + + public HttpMessageHandler CreateMessageHandler() + { + Mock handler = new() { CallBase = true }; + handler.Protected() + .Setup>("SendAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(this.SendAsync); + if (this.FailTransportDisposal) + { + handler.Protected().Setup("Dispose", ItExpr.Is(disposing => disposing)) + .Throws(new InvalidOperationException("transport cleanup failed")); + } + + this.Handlers.Add(handler); + return handler.Object; + } + + private async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + if (request.Method == HttpMethod.Get) + { + return EmptyResponse(HttpStatusCode.MethodNotAllowed, request); + } + + if (request.Method == HttpMethod.Delete) + { + Interlocked.Increment(ref this._terminations); + return EmptyResponse(HttpStatusCode.OK, request); + } + +#if NET + using JsonDocument document = JsonDocument.Parse(await request.Content!.ReadAsStringAsync(cancellationToken)); +#else + using JsonDocument document = JsonDocument.Parse(await request.Content!.ReadAsStringAsync()); +#endif + JsonElement message = document.RootElement; + if (!message.TryGetProperty("id", out JsonElement id)) + { + return EmptyResponse(HttpStatusCode.Accepted, request); + } + + string? method = message.GetProperty("method").GetString(); + if (method == "server/discover") + { + return new(HttpStatusCode.OK) + { + RequestMessage = request, + Content = new StringContent( + $$$"""{"jsonrpc":"2.0","id":{{{id.GetRawText()}}},"error":{"code":-32601,"message":"Method not found"}}""", + Encoding.UTF8, "application/json") + }; + } + + string result; + string? sessionId = null; + if (method == "initialize") + { + if (this.FailInitialization) + { + return EmptyResponse(HttpStatusCode.BadRequest, request); + } + + sessionId = Interlocked.Increment(ref this._initializations).ToString(CultureInfo.InvariantCulture); + string protocolVersion = message.GetProperty("params").GetProperty("protocolVersion").GetRawText(); + result = $$$"""{"protocolVersion":{{{protocolVersion}}},"capabilities":{"tools":{}},"serverInfo":{"name":"stub","version":"1.0"}}"""; + } + else + { + if (this.BeforeOperationAsync is not null) + { + await this.BeforeOperationAsync(cancellationToken); + } + + if (this.FailOperation) + { + return EmptyResponse(HttpStatusCode.BadRequest, request); + } + + result = method switch + { + "tools/list" => """{"tools":[]}""", + "tools/call" => """{"content":[{"type":"text","text":"ok"}]}""", + _ => throw new InvalidOperationException($"Unexpected MCP method: {method}") + }; + } + + HttpResponseMessage response = new(HttpStatusCode.OK) + { + RequestMessage = request, + Content = new StringContent( + $$$"""{"jsonrpc":"2.0","id":{{{id.GetRawText()}}},"result":{{{result}}}}""", + Encoding.UTF8, "application/json") + }; + if (sessionId is not null) + { + response.Headers.Add("Mcp-Session-Id", sessionId); + } + + return response; + } + + private static HttpResponseMessage EmptyResponse(HttpStatusCode statusCode, HttpRequestMessage request) => + new(statusCode) + { + RequestMessage = request, + Content = new ByteArrayContent([]) + }; + } + + private sealed class CleanupTraceListener : TraceListener + { + private readonly StringBuilder _output = new(); + + public string Output => this._output.ToString(); + + public override void Write(string? message) => this._output.Append(message); + + public override void WriteLine(string? message) => this._output.AppendLine(message); + } +} diff --git a/python/packages/declarative/AGENTS.md b/python/packages/declarative/AGENTS.md index 3c9402fc4e5..17229ecaa9c 100644 --- a/python/packages/declarative/AGENTS.md +++ b/python/packages/declarative/AGENTS.md @@ -12,6 +12,21 @@ YAML/JSON-based declarative agent and workflow definitions. - **`MCPToolHandler`** / **`DefaultMCPToolHandler`** - Pluggable MCP transport for the `InvokeMcpTool` declarative action (configured via `WorkflowFactory(mcp_tool_handler=...)`) - **`DeclarativeLoaderError`** / **`ProviderLookupError`** / **`DeclarativeWorkflowError`** / **`DeclarativeActionError`** - Error types +## MCP Handler Lifetimes + +`DefaultMCPToolHandler` caches/coalesces sessions only without a `client_provider`. +With a provider, every invocation (including `tools/list`) gets a fresh tool/session, +even if the provider returns `None` or a shared HTTP client. Invocation cleanup closes +the session and any internally owned fallback client, never caller-owned HTTP clients. +Shutdown waits for active provider-backed invocations to clean up. Calling `aclose()` +from an active invocation's context (including inherited child tasks and cleanup) +raises `RuntimeError` before changing handler state; shutdown must run outside that +context or after invocation completion. + +Per-invocation sessions intentionally add connection overhead and lose server +session continuity between invocations; +shared session ownership requires an explicitly scoped custom `MCPToolHandler`. + ## External Input Handling - **`ExternalInputRequest`** / **`ExternalInputResponse`** - Human-in-the-loop support diff --git a/python/packages/declarative/README.md b/python/packages/declarative/README.md index c536403aba4..42a2a2bc305 100644 --- a/python/packages/declarative/README.md +++ b/python/packages/declarative/README.md @@ -39,3 +39,25 @@ owned client with the default cookie policy. Explicit `Cookie` request headers remain supported. Response `Set-Cookie` headers remain available in `HttpRequestResult.headers`; disabling persistence does not redact the response. + +## MCP handler session lifetime + +`DefaultMCPToolHandler()` reuses MCP sessions through a bounded LRU cache. +When configured with `client_provider`, it instead invokes the provider and +creates a fresh MCP tool/session for every invocation, including `tools/list`. +Sessions are never reused in this mode, even if the provider returns `None` +or the same `httpx.AsyncClient` instance. + +Each invocation closes its tool/session on success, failure, or cancellation. +Internally created fallback HTTP clients are also closed; caller-supplied HTTP +clients remain caller-owned and are never closed by the handler. Calling +`aclose()` rejects new invocations and waits for active provider-backed +invocations to finish cleaning up. Calling it from an active invocation's +context (including provider callbacks, inherited child tasks, and cleanup) +raises `RuntimeError` before changing handler state, rather than waiting on +itself. Close the handler outside that context or after the invocation completes. + +This intentionally incurs connection and initialization overhead and does not +retain server session state between provider-backed invocations. Applications +that require shared session ownership must implement an explicitly scoped +custom `MCPToolHandler`; there is no provider-backed session-cache opt-in. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py index 4a26e74c885..64eb59750fb 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_mcp_handler.py @@ -32,6 +32,7 @@ import logging from collections import OrderedDict from collections.abc import Awaitable, Callable +from contextvars import ContextVar, Token from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Protocol, cast, runtime_checkable @@ -157,14 +158,14 @@ class _CacheEntry: class DefaultMCPToolHandler: """Default :class:`MCPToolHandler` backed by :class:`agent_framework.MCPStreamableHTTPTool`. - Caches one :class:`agent_framework.MCPStreamableHTTPTool` instance per - ``(server_url, server_label, connection_name, headers_hash)`` in a - bounded LRU. The cache prevents re-establishing an MCP session for every + Without a ``client_provider``, caches one + :class:`agent_framework.MCPStreamableHTTPTool` instance per + ``(server_url, server_label, connection_name, headers_hash)`` in a bounded + LRU. The cache prevents re-establishing an MCP session for every invocation while ensuring different header sets (auth tokens) cannot share a session — matches the .NET design intent while bounding - cardinality. ``server_label`` and ``connection_name`` participate in - the key so that callers using ``client_provider`` to dispatch on those - fields receive a fresh client per logical connection (see below). + cardinality. ``server_label`` and ``connection_name`` also participate + in the key to distinguish logical connections. Header *names* are lower-cased inside the hash payload only — the headers passed on the wire keep the caller's original casing — so two YAML actions that spell ``Authorization`` differently still share a @@ -174,12 +175,20 @@ class DefaultMCPToolHandler: 1. ``DefaultMCPToolHandler()`` — owns its own ``httpx.AsyncClient`` instances created lazily per cache entry. Closed by :meth:`aclose`. - 2. ``DefaultMCPToolHandler(client_provider=cb)`` — per-server client + 2. ``DefaultMCPToolHandler(client_provider=cb)`` — per-invocation client lookup (parity with .NET ``httpClientProvider`` callback). The callback receives the full :class:`MCPToolInvocation` so it can dispatch on ``server_url`` / ``connection_name`` / ``server_label``. - Returning ``None`` falls back to an internally-created client. Caller - supplied clients are NOT closed by :meth:`aclose`. + The callback is invoked for every call, including ``tools/list``. + Each call creates and closes its own MCP tool/session, even when the + callback returns ``None`` or the same client object. Returning ``None`` + falls back to an internally-created client, closed with the invocation. + Caller-supplied clients are never closed by this handler. + + This intentionally trades connection/session reuse for invocation + isolation. Server session state is not retained across invocations; + callers needing shared session ownership must implement an explicitly + scoped custom :class:`MCPToolHandler`. .. warning:: @@ -187,11 +196,11 @@ class DefaultMCPToolHandler: or replace it with a custom handler in production deployments. Args: - client_provider: Optional per-server ``httpx.AsyncClient`` provider. - cache_max_size: Maximum number of cached MCP clients. When exceeded, - the least-recently-used entry is evicted and its client closed - (only owned clients are closed; caller-supplied ones are not). - Defaults to ``32``. + client_provider: Optional per-invocation ``httpx.AsyncClient`` provider. + cache_max_size: Maximum number of cached MCP clients in no-provider mode. + When exceeded, the least-recently-used entry is evicted and its + owned client closed. Defaults to ``32``. Does not enable session + caching when a provider is configured. """ LIST_TOOLS_TOOL_NAME: ClassVar[str] = "tools/list" @@ -227,12 +236,19 @@ def __init__( # tasks awaiting the same key will await the same future and share # the resulting cache entry. self._inflight: dict[tuple[str, str, str, str], asyncio.Future[_CacheEntry]] = {} + # Completion signals only: provider-backed calls never share entries. + self._active_invocations: set[asyncio.Future[None]] = set() + # Keep ancestry so a completed nested call cannot hide an active parent + # in the context inherited by child tasks, including cleanup tasks. + self._invocation_context: ContextVar[tuple[asyncio.Future[None], ...]] = ContextVar( + f"default_mcp_tool_handler_invocations_{id(self)}", default=() + ) # Set by ``aclose`` to prevent post-close cache insertions and to # reject new ``invoke_tool`` calls. Once set, never cleared. self._closed = False async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: - """Invoke ``invocation.tool_name`` on the cached MCP client for the server. + """Invoke ``invocation.tool_name`` on an MCP client for the server. The reserved name :attr:`LIST_TOOLS_TOOL_NAME` (``"tools/list"``) is intercepted client-side: instead of being forwarded as a tool call, @@ -241,7 +257,6 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: ``TextContent`` containing a JSON tool catalog. """ from agent_framework import Content - from agent_framework.exceptions import ToolExecutionException # Reserved-name args validation runs before connect: rejecting bad # input shouldn't require establishing an MCP session. @@ -253,23 +268,55 @@ async def invoke_tool(self, invocation: MCPToolInvocation) -> MCPToolResult: error_message=message, ) + entry: _CacheEntry | None = None + completion: asyncio.Future[None] | None = None + context_token: Token[tuple[asyncio.Future[None], ...]] | None = None try: - entry = await self._get_or_create_entry(invocation) - except Exception as exc: - # Connect / cache lookup failures surface as tool errors so the - # workflow can store them at output.result without crashing. - logger.warning( - "DefaultMCPToolHandler: failed to obtain MCP client for url=%s tool=%s: %s", - invocation.server_url, - invocation.tool_name, - exc, - ) - message = f"Failed to connect to MCP server: {type(exc).__name__}: {exc}".rstrip(": ") - return MCPToolResult( - outputs=[Content.from_text(f"Error: {message}")], - is_error=True, - error_message=message, - ) + try: + if self._client_provider is None: + entry = await self._get_or_create_entry(invocation) + else: + async with self._cache_lock: + if self._closed: + raise RuntimeError("DefaultMCPToolHandler is closed") + completion = asyncio.get_running_loop().create_future() + self._active_invocations.add(completion) + context_token = self._invocation_context.set((*self._invocation_context.get(), completion)) + entry = await self._create_entry(invocation) + if self._closed: + raise RuntimeError("DefaultMCPToolHandler is closed") + except Exception as exc: + # Connect / cache lookup failures surface as tool errors so the + # workflow can store them at output.result without crashing. + logger.warning( + "DefaultMCPToolHandler: failed to obtain MCP client for url=%s tool=%s: %s", + invocation.server_url, + invocation.tool_name, + exc, + ) + message = f"Failed to connect to MCP server: {type(exc).__name__}: {exc}".rstrip(": ") + return MCPToolResult( + outputs=[Content.from_text(f"Error: {message}")], + is_error=True, + error_message=message, + ) + + return await self._invoke_entry(entry, invocation) + finally: + if completion is not None: + try: + if entry is not None: + await self._close_invocation_entry(entry) + finally: + if context_token is not None: + self._invocation_context.reset(context_token) + self._active_invocations.discard(completion) + completion.set_result(None) + + async def _invoke_entry(self, entry: _CacheEntry, invocation: MCPToolInvocation) -> MCPToolResult: + """Dispatch a connected invocation and preserve tool error mapping.""" + from agent_framework import Content + from agent_framework.exceptions import ToolExecutionException try: if invocation.tool_name == self.LIST_TOOLS_TOOL_NAME: @@ -372,24 +419,42 @@ async def _invoke_list_tools(entry: _CacheEntry) -> MCPToolResult: return MCPToolResult(outputs=[Content.from_text(json.dumps(payload, indent=2, allow_nan=False))]) async def aclose(self) -> None: - """Close all cached MCP clients and the owned httpx clients. + """Close cached clients and wait for provider-backed invocations to clean up. Caller-supplied :class:`httpx.AsyncClient` instances (returned by the ``client_provider`` callback) are NOT closed. - Idempotent — a second call returns immediately. Drains any in-flight - ``_create_entry`` tasks before returning so their resources are + Provider-backed calls already executing may finish; pending connections + are rejected after connecting and cleaned up by their invocation. + Concurrent shutdown calls wait for those invocations as well. + + Idempotent. In no-provider mode, a second call returns immediately. + Drains any in-flight ``_create_entry`` tasks before returning so their resources are cleaned up; the in-flight tasks see ``self._closed`` in phase 3 of :meth:`_get_or_create_entry`, close their own entry, and resolve their future with ``RuntimeError("DefaultMCPToolHandler is closed")``. + + Raises: + RuntimeError: If called from an active provider-backed invocation's + context, including inherited child tasks and cleanup. Rejected + before changing handler state to avoid waiting on itself. """ + if any(not completion.done() for completion in self._invocation_context.get()): + raise RuntimeError( + "DefaultMCPToolHandler.aclose() cannot be called from an active provider-backed invocation" + ) async with self._cache_lock: - if self._closed: + if self._closed and self._client_provider is None: return self._closed = True entries = list(self._cache.values()) self._cache.clear() inflight_futures = list(self._inflight.values()) + active_invocations = list(self._active_invocations) + + for completion in active_invocations: + # Cancelling shutdown must not cancel an invocation's completion signal. + await asyncio.shield(completion) # Wait for in-flight creations to finish their self-cleanup. Each # in-flight task self-closes its entry under the closed-flag branch @@ -509,7 +574,9 @@ async def _create_entry(self, invocation: MCPToolInvocation) -> _CacheEntry: provided_client: httpx.AsyncClient | None = None if self._client_provider is not None: provided_client = await self._client_provider(invocation) - # Capture headers for this cache entry so the header_provider closure + if self._closed: + raise RuntimeError("DefaultMCPToolHandler is closed") + # Capture headers for this entry so the header_provider closure # always returns the same set, regardless of the runtime kwargs. captured_headers = dict(invocation.headers) @@ -526,10 +593,18 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: try: await tool.connect() except BaseException: - try: - await tool.close() - except Exception: # pragma: no cover - best effort - logger.debug("DefaultMCPToolHandler: error closing tool after failed connect", exc_info=True) + failed_entry = _CacheEntry( + tool=tool, + owned_httpx_client=( + cast("httpx.AsyncClient | None", getattr(tool, "_httpx_client", None)) + if provided_client is None + else None + ), + ) + if self._client_provider is not None: + await self._close_invocation_entry(failed_entry) + else: + await self._close_entry(failed_entry) raise # ``MCPStreamableHTTPTool.get_mcp_client`` lazily creates an @@ -542,17 +617,34 @@ def _header_provider(_kwargs: dict[str, Any]) -> dict[str, str]: owned_client = cast("httpx.AsyncClient | None", getattr(tool, "_httpx_client", None)) return _CacheEntry(tool=tool, owned_httpx_client=owned_client) + async def _close_invocation_entry(self, entry: _CacheEntry) -> None: + """Finish invocation cleanup even if the caller is cancelled again.""" + # MCPStreamableHTTPTool dispatches connect/close to its lifecycle owner, + # keeping the SDK's cancel-scope entry and exit on that same task. + cleanup = asyncio.create_task(self._close_entry(entry)) + cancelled = False + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + cancelled = True + cleanup.result() # Propagate cancellation/errors from cleanup itself. + if cancelled: + raise asyncio.CancelledError + async def _close_entry(self, entry: _CacheEntry) -> None: """Close the MCP tool and any owned httpx client.""" try: - await entry.tool.close() - except Exception: # pragma: no cover - best effort - logger.debug("DefaultMCPToolHandler: error closing MCP tool", exc_info=True) - if entry.owned_httpx_client is not None: try: - await entry.owned_httpx_client.aclose() + await entry.tool.close() except Exception: # pragma: no cover - best effort - logger.debug("DefaultMCPToolHandler: error closing owned httpx client", exc_info=True) + logger.debug("DefaultMCPToolHandler: error closing MCP tool", exc_info=True) + finally: + if entry.owned_httpx_client is not None and not entry.owned_httpx_client.is_closed: + try: + await entry.owned_httpx_client.aclose() + except Exception: # pragma: no cover - best effort + logger.debug("DefaultMCPToolHandler: error closing owned httpx client", exc_info=True) @staticmethod def _cache_key( @@ -563,10 +655,9 @@ def _cache_key( ) -> tuple[str, str, str, str]: """Build an order-independent cache key for the invocation identity. - The key includes ``server_label`` and ``connection_name`` so that - callers using ``client_provider`` to dispatch on those fields - receive a fresh client per logical connection (matches the - documented dispatch contract). + Used only without a ``client_provider``. The key includes + ``server_label`` and ``connection_name`` to distinguish logical + connections. Header *names* are lower-cased inside the hash payload only so that ``Authorization`` and ``authorization`` map to the same diff --git a/python/packages/declarative/tests/test_default_mcp_tool_handler.py b/python/packages/declarative/tests/test_default_mcp_tool_handler.py index a8329dea278..0826302cf69 100644 --- a/python/packages/declarative/tests/test_default_mcp_tool_handler.py +++ b/python/packages/declarative/tests/test_default_mcp_tool_handler.py @@ -7,10 +7,9 @@ exercisable through the executor stub: cache hit/miss/eviction, concurrent connect via in-flight futures, header isolation across cache keys, string-result normalisation, ``load_prompts=False`` verification, and -owned-vs-caller httpx close semantics. +owned-vs-caller httpx close semantics, and provider-backed invocation lifetimes. -The shared-client regression also exercises the real MCP SDK transport against -an in-process HTTPX mock server. +Two lifecycle cases use the real MCP SDK with an inert HTTPX mock transport. """ from __future__ import annotations @@ -19,7 +18,7 @@ import json import sys from typing import Any -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import httpx import pytest @@ -102,7 +101,7 @@ def __init__(self, **kwargs: Any) -> None: self.connect_delay: float = 0.0 self.connect_error: BaseException | None = None self.call_handler: Any = lambda **_a: [Content.from_text("ok")] - self._httpx_client: httpx.AsyncClient | None = None + self._httpx_client: httpx.AsyncClient | None = kwargs.get("http_client") self.session: FakeMcpSession | None = None # Mimic MCPStreamableHTTPTool: when no caller client AND header_provider # is set, lazily allocate an owned httpx client during connect. @@ -150,81 +149,671 @@ def _invocation( ) -@pytest.mark.parametrize("cache_max_size", [1, 2]) -async def test_shared_client_isolates_cached_authentication_and_cleans_up_hooks(cache_max_size: int) -> None: - sessions: dict[str, str] = {} - calls: list[tuple[str, str]] = [] - writes: dict[str, list[str]] = {"token-a": [], "token-b": []} - - async def caller_hook(request: httpx.Request) -> None: - request.headers["X-Caller"] = "preserved" - - async def handle(request: httpx.Request) -> httpx.Response: - assert request.headers["X-Caller"] == "preserved" - principal = request.headers.get("Authorization", "") - if principal not in writes: - return httpx.Response(401) - if request.method == "GET": - return httpx.Response(405) - if request.method == "DELETE": - return httpx.Response(200) - body = json.loads(request.content) - headers: dict[str, str] = {} - result: dict[str, Any] = {} - if body.get("method") == "initialize": - session_id = f"session-{len(sessions)}" - sessions[session_id] = principal - headers["mcp-session-id"] = session_id - result = { - "protocolVersion": body["params"]["protocolVersion"], - "capabilities": {"tools": {}}, - "serverInfo": {"name": "auth-test", "version": "1"}, - } - elif body.get("method") == "tools/list": - result = { - "tools": [ - { - "name": "search", - "inputSchema": {"type": "object", "properties": {"marker": {"type": "string"}}}, - } - ] - } - elif body.get("method") == "tools/call": - calls.append((request.headers["mcp-session-id"], principal)) - if marker := body["params"].get("arguments", {}).get("marker"): - writes[principal].append(marker) - result = {"content": [{"type": "text", "text": principal}]} - if "id" not in body: - return httpx.Response(202) - return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) - - async with httpx.AsyncClient( - transport=httpx.MockTransport(handle), event_hooks={"request": [caller_hook]} - ) as client: - - async def client_provider(invocation: MCPToolInvocation) -> httpx.AsyncClient: - return client - - async with DefaultMCPToolHandler(client_provider=client_provider, cache_max_size=cache_max_size) as handler: - outputs: list[str | None] = [] - for index, principal in enumerate(("token-a", "token-b", "token-a")): - result = await handler.invoke_tool( - _invocation( - headers={"Authorization": principal}, - arguments={"marker": "a-only"} if index == 2 else {}, - ) +# ---------- Provider-backed invocation lifetimes -------------------------- + + +class TestProviderLifetimes: + @pytest.mark.parametrize("returns_none", [False, True]) + async def test_actual_sdk_lifecycle_and_hook_cleanup( + self, returns_none: bool, caplog: pytest.LogCaptureFixture + ) -> None: + import agent_framework + from agent_framework import MCPStreamableHTTPTool + + caplog.set_level("DEBUG", logger="agent_framework_declarative._workflows._mcp_handler") + initialized: list[str] = [] + terminated: list[str] = [] + methods: list[str] = [] + tools: list[MCPStreamableHTTPTool] = [] + clients: list[RecordingClient] = [] + caller_hook = AsyncMock() + + async def respond(request: httpx.Request) -> httpx.Response: + assert request.headers["X-Test"] == "lifecycle" + if request.method == "GET": + return httpx.Response(405) + if request.method == "DELETE": + terminated.append(request.headers["mcp-session-id"]) + return httpx.Response(200) + + body = json.loads(request.content) + method = body["method"] + methods.append(method) + if "id" not in body: + return httpx.Response(202) + headers: dict[str, str] = {} + if method == "initialize": + session_id = f"lifecycle-{len(initialized) + 1}" + initialized.append(session_id) + headers["mcp-session-id"] = session_id + result = { + "protocolVersion": body["params"]["protocolVersion"], + "capabilities": {"tools": {}}, + "serverInfo": {"name": "lifecycle", "version": "1"}, + } + elif method == "tools/list": + result = { + "tools": [{"name": "ok", "inputSchema": {"type": "object", "properties": {}}}], + } + elif method == "tools/call": + result = {"content": [{"type": "text", "text": "ok"}]} + else: + assert method == "ping" + result = {} + return httpx.Response(200, headers=headers, json={"jsonrpc": "2.0", "id": body["id"], "result": result}) + + class RecordingClient(httpx.AsyncClient): + def __init__(self, **kwargs: Any) -> None: + super().__init__( + transport=httpx.MockTransport(respond), event_hooks={"request": [caller_hook]}, **kwargs ) + self.close_count = 0 + clients.append(self) + + async def aclose(self) -> None: + self.close_count += 1 + await super().aclose() + + def create_tool(**kwargs: Any) -> MCPStreamableHTTPTool: + tool = MCPStreamableHTTPTool(**kwargs) + tools.append(tool) + return tool + + caller_client = None if returns_none else RecordingClient() + provider = AsyncMock(return_value=caller_client) + try: + with ( + patch.object(httpx, "AsyncClient", RecordingClient), + patch.object(agent_framework, "MCPStreamableHTTPTool", side_effect=create_tool), + ): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + for index, tool_name in enumerate(("ok", "tools/list", "ok", "tools/list"), start=1): + result = await handler.invoke_tool( + _invocation(tool_name=tool_name, headers={"X-Test": "lifecycle"}) + ) + assert not result.is_error + if tool_name == "ok": + assert result.outputs[0].text == "ok" + else: + assert result.outputs[0].text is not None + assert json.loads(result.outputs[0].text)["tools"][0]["name"] == "ok" + assert provider.await_count == index + assert len(tools) == index + assert len(initialized) == index + assert terminated == initialized + assert not handler._active_invocations + assert not handler._cache + tool = tools[-1] + assert tool.session is None + assert not tool.is_connected + assert tool._lifecycle_owner_task is None + assert tool._header_hook_client is None + assert all(client.event_hooks["request"] == [caller_hook] for client in clients) + if returns_none: + assert len(clients) == index + assert all(client.is_closed and client.close_count == 1 for client in clients) + else: + assert caller_client is not None + assert not caller_client.is_closed + assert caller_client.close_count == 0 + assert len({id(tool) for tool in tools}) == 4 + assert methods.count("tools/call") == 2 + assert methods.count("tools/list") == 6 + assert not any( + "Could not cleanly close MCP exit stack" in record.getMessage() + or "DefaultMCPToolHandler: error closing" in record.getMessage() + for record in caplog.records + ) + finally: + for client in clients: + if not client.is_closed: + await client.aclose() + + @pytest.mark.parametrize("returns_none", [False, True]) + async def test_sequential_invocations_get_fresh_tools(self, returns_none: bool) -> None: + async with httpx.AsyncClient() as client: + provider = AsyncMock(return_value=None if returns_none else client) + with _patch_tool(): + async with DefaultMCPToolHandler(client_provider=provider, cache_max_size=1) as handler: + for index, tool_name in enumerate(("search", "tools/list", "search", "tools/list"), start=1): + invocation = _invocation(tool_name=tool_name, headers={"X-Test": "1"}) + result = await handler.invoke_tool(invocation) + assert not result.is_error + assert result.outputs[0].text == ('{\n "tools": []\n}' if tool_name == "tools/list" else "ok") + assert provider.await_count == index + assert provider.await_args is not None + assert provider.await_args.args[0] is invocation + assert len(FakeTool.instances) == index + tool = FakeTool.instances[-1] + assert tool.connect_count == 1 + assert tool.close_count == 1 + assert tool.kwargs["http_client"] is (None if returns_none else client) + if returns_none: + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + assert not client.is_closed + assert not handler._cache + assert not handler._inflight + assert not handler._active_invocations + assert len({id(tool.session) for tool in FakeTool.instances}) == 4 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + assert not client.is_closed + + @pytest.mark.parametrize("returns_none", [False, True]) + async def test_concurrent_invocations_get_fresh_tools(self, returns_none: bool) -> None: + all_connecting = asyncio.Event() + release = asyncio.Event() + original_connect = FakeTool.connect + + async def gated_connect(tool: FakeTool) -> None: + await original_connect(tool) + if len(FakeTool.instances) == 3: + all_connecting.set() + await release.wait() + + async with httpx.AsyncClient() as client: + provider = AsyncMock(return_value=None if returns_none else client) + with _patch_tool(), patch.object(FakeTool, "connect", gated_connect): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + tasks = [ + asyncio.create_task(handler.invoke_tool(_invocation(tool_name=name, headers={"X-Test": "1"}))) + for name in ("search", "tools/list", "search") + ] + try: + await asyncio.wait_for(all_connecting.wait(), timeout=5) + assert provider.await_count == 3 + assert len({id(tool.session) for tool in FakeTool.instances}) == 3 + assert all(tool.close_count == 0 for tool in FakeTool.instances) + assert not handler._cache + assert not handler._inflight + finally: + release.set() + results = await asyncio.gather(*tasks) + assert all(not result.is_error for result in results) + assert all(tool.close_count == 1 for tool in FakeTool.instances) + if returns_none: + assert all(tool._httpx_client and tool._httpx_client.is_closed for tool in FakeTool.instances) + assert not handler._active_invocations + assert not client.is_closed + + @pytest.mark.timeout(10) + @pytest.mark.parametrize("stage", ["provider", "connect", "search", "tools/list", "close"]) + @pytest.mark.parametrize("in_child_task", [False, True]) + async def test_reentrant_shutdown_is_rejected_without_closing_handler( + self, stage: str, in_child_task: bool + ) -> None: + original_connect = FakeTool.connect + original_close = FakeTool.close + original_call = FakeTool.call_tool + original_list = FakeMcpSession.list_tools + rejected = 0 + + async def check_shutdown(phase: str) -> None: + nonlocal rejected + if stage != phase: + return + with pytest.raises(RuntimeError, match="cannot be called from an active provider-backed invocation"): + if in_child_task: + await asyncio.create_task(handler.aclose()) + else: + await handler.aclose() + rejected += 1 + assert not handler._closed + + async def provider(_invocation: MCPToolInvocation) -> None: + await check_shutdown("provider") + + async def connect(tool: FakeTool) -> None: + await original_connect(tool) + await check_shutdown("connect") + + async def close(tool: FakeTool) -> None: + await check_shutdown("close") + await original_close(tool) + + async def call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + await check_shutdown("search") + return await original_call(tool, tool_name, **arguments) + + async def list_tools(session: FakeMcpSession, params: Any = None) -> FakeListToolsResult: + await check_shutdown("tools/list") + return await original_list(session, params) + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", connect), + patch.object(FakeTool, "close", close), + patch.object(FakeTool, "call_tool", call), + patch.object(FakeMcpSession, "list_tools", list_tools), + ): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + for _ in range(2): + result = await handler.invoke_tool( + _invocation(tool_name="tools/list" if stage == "tools/list" else "search") + ) + assert not result.is_error + assert handler._invocation_context.get() == () + assert not handler._active_invocations + assert not handler._closed + assert rejected == 2 + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.timeout(10) + async def test_nested_invocations_restore_context_and_keep_active_ancestry(self) -> None: + release_child = asyncio.Event() + children: list[asyncio.Task[None]] = [] + + async def close_from_inherited_context() -> None: + await release_child.wait() + outer, inner = handler._invocation_context.get() + assert not outer.done() + assert inner.done() + with pytest.raises(RuntimeError, match="cannot be called from an active provider-backed invocation"): + await handler.aclose() + assert not handler._closed + + async def provider(invocation: MCPToolInvocation) -> None: + if invocation.tool_name == "inner": + children.append(asyncio.create_task(close_from_inherited_context())) + return + outer_context = handler._invocation_context.get() + assert len(outer_context) == 1 + nested_result = await handler.invoke_tool(_invocation(tool_name="inner")) + assert not nested_result.is_error + assert handler._invocation_context.get() is outer_context + release_child.set() + await children[-1] + with pytest.raises(RuntimeError, match="cannot be called from an active provider-backed invocation"): + await handler.aclose() + + with _patch_tool(): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + try: + result = await handler.invoke_tool(_invocation(tool_name="outer")) + assert not result.is_error + assert handler._invocation_context.get() == () + assert not handler._active_invocations + finally: + release_child.set() + await asyncio.gather(*children) + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.timeout(10) + async def test_provider_can_invoke_and_close_another_handler(self) -> None: + async def provider(_invocation: MCPToolInvocation) -> None: + async with DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) as other: + result = await other.invoke_tool(_invocation) + assert not result.is_error + assert other._invocation_context.get() == () + assert handler._invocation_context.get() + assert other._closed + assert not handler._closed + + with _patch_tool(): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + result = await handler.invoke_tool(_invocation()) + assert not result.is_error + assert handler._invocation_context.get() == () + assert len(FakeTool.instances) == 2 + assert all(tool.close_count == 1 for tool in FakeTool.instances) + + @pytest.mark.timeout(10) + async def test_inherited_context_can_close_after_invocation_completes(self) -> None: + release_child = asyncio.Event() + children: list[asyncio.Task[None]] = [] + + async def close_from_inherited_context() -> None: + await release_child.wait() + context = handler._invocation_context.get() + assert len(context) == 1 + assert context[0].done() + await handler.aclose() + + async def provider(_invocation: MCPToolInvocation) -> None: + children.append(asyncio.create_task(close_from_inherited_context())) + + with _patch_tool(): + handler = DefaultMCPToolHandler(client_provider=provider) + try: + result = await handler.invoke_tool(_invocation()) + assert not result.is_error + assert handler._invocation_context.get() == () + assert not handler._active_invocations + assert not handler._closed + finally: + release_child.set() + await asyncio.gather(*children) + await handler.aclose() + assert handler._closed + assert len(FakeTool.instances) == 1 + assert FakeTool.instances[0].close_count == 1 + + async def test_list_tools_paginates_within_each_invocation(self) -> None: + original_connect = FakeTool.connect + + async def connect(tool: FakeTool) -> None: + await original_connect(tool) + assert tool.session is not None + tool.session.list_tools_pages = [ + FakeListToolsResult([FakeMcpTool("first")], next_cursor="next"), + FakeListToolsResult([FakeMcpTool("second")]), + ] + + provider = AsyncMock(return_value=None) + with ( + _patch_tool(), + patch.object(FakeTool, "connect", connect), + patch.object(FakeTool, "call_tool", new_callable=AsyncMock) as call_tool, + ): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + for _ in range(2): + result = await handler.invoke_tool(_invocation(tool_name="tools/list")) + assert not result.is_error + assert result.outputs[0].text is not None + payload = json.loads(result.outputs[0].text) + assert [tool["name"] for tool in payload["tools"]] == ["first", "second"] + call_tool.assert_not_awaited() + assert provider.await_count == 2 + assert len(FakeTool.instances) == 2 + for tool in FakeTool.instances: + assert tool.close_count == 1 + assert tool.session is not None + assert len(tool.session.list_tools_calls) == 2 + assert tool.session.list_tools_calls[0] is None + assert tool.session.list_tools_calls[1].cursor == "next" + + @pytest.mark.parametrize("stage", ["provider", "connect", "search", "tools/list"]) + @pytest.mark.parametrize( + "error_type", [httpx.ReadTimeout, ToolExecutionException, RuntimeError, asyncio.CancelledError] + ) + @pytest.mark.parametrize("returns_none", [False, True]) + async def test_errors_clean_up_and_preserve_mapping( + self, stage: str, error_type: type[BaseException], returns_none: bool + ) -> None: + error = error_type("operation stopped") + original_connect = FakeTool.connect + + async def connect(tool: FakeTool) -> None: + await original_connect(tool) + if stage == "connect": + raise error + + async with httpx.AsyncClient() as client: + provider = AsyncMock( + return_value=None if returns_none else client, + side_effect=error if stage == "provider" else None, + ) + with ( + _patch_tool(), + patch.object(FakeTool, "connect", connect), + patch.object(FakeTool, "call_tool", new_callable=AsyncMock, side_effect=error), + patch.object(FakeMcpSession, "list_tools", new_callable=AsyncMock, side_effect=error), + ): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + invocation = _invocation( + tool_name="tools/list" if stage == "tools/list" else "search", headers={"X-Test": "1"} + ) + if error_type is asyncio.CancelledError or ( + error_type is RuntimeError and stage in ("search", "tools/list") + ): + with pytest.raises(error_type, match="operation stopped"): + await handler.invoke_tool(invocation) + else: + result = await handler.invoke_tool(invocation) + assert result.is_error + assert "operation stopped" in (result.error_message or "") + assert result.outputs[0].text is not None + assert result.outputs[0].text.startswith("Error:") + assert handler._invocation_context.get() == () + assert not handler._active_invocations + assert not handler._cache + assert not handler._inflight + assert provider.await_count == 1 + assert len(FakeTool.instances) == (0 if stage == "provider" else 1) + for tool in FakeTool.instances: + assert tool.close_count == 1 + if returns_none: + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + assert not client.is_closed + + @pytest.mark.parametrize("tool_name", ["search", "tools/list"]) + async def test_mcp_error_mapping_cleans_up(self, tool_name: str) -> None: + from mcp.shared.exceptions import McpError + from mcp.types import ErrorData + + error = McpError(ErrorData(code=-32603, message="operation stopped")) + with ( + _patch_tool(), + patch.object(FakeTool, "call_tool", new_callable=AsyncMock, side_effect=error), + patch.object(FakeMcpSession, "list_tools", new_callable=AsyncMock, side_effect=error), + ): + async with DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) as handler: + result = await handler.invoke_tool(_invocation(tool_name=tool_name, headers={"X-Test": "1"})) + assert result.is_error + assert result.error_message == "operation stopped" + assert not handler._active_invocations + tool = FakeTool.instances[0] + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + + async def test_tool_close_failure_still_closes_fallback_client(self) -> None: + async def close(tool: FakeTool) -> None: + tool.close_count += 1 + raise RuntimeError("close failed") + + with _patch_tool(), patch.object(FakeTool, "close", close): + async with DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) as handler: + result = await handler.invoke_tool(_invocation(headers={"X-Test": "1"})) assert not result.is_error - outputs.append(result.outputs[0].text) - assert outputs == ["token-a", "token-b", "token-a"] - assert len(client.event_hooks["request"]) == 1 + cache_max_size - - assert [principal for _, principal in calls] == ["token-a", "token-b", "token-a"] - assert all(sessions[session_id] == principal for session_id, principal in calls) - assert len(sessions) == (3 if cache_max_size == 1 else 2) - assert writes == {"token-a": ["a-only"], "token-b": []} - assert client.event_hooks["request"] == [caller_hook] - assert not client.is_closed + assert result.outputs[0].text == "ok" + assert not handler._active_invocations + tool = FakeTool.instances[0] + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + + @pytest.mark.parametrize("connect_fails", [False, True]) + async def test_cleanup_cancellation_signals_completion(self, connect_fails: bool) -> None: + original_connect = FakeTool.connect + completions: list[asyncio.Future[None]] = [] + + async def connect(tool: FakeTool) -> None: + await original_connect(tool) + if connect_fails: + raise httpx.ConnectError("connect stopped") + + async def close(tool: FakeTool) -> None: + tool.close_count += 1 + completions.extend(handler._active_invocations) + raise asyncio.CancelledError("cleanup stopped") + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", connect), + patch.object(FakeTool, "close", close), + ): + handler = DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) + # Python 3.10 may drop the message when a cancelled task's result is retrieved again. + with pytest.raises(asyncio.CancelledError): + await handler.invoke_tool(_invocation(headers={"X-Test": "1"})) + assert handler._invocation_context.get() == () + assert not handler._active_invocations + assert len(completions) == 1 + assert completions[0].done() + assert completions[0].result() is None + await asyncio.wait_for(handler.aclose(), timeout=5) + tool = FakeTool.instances[0] + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + + @pytest.mark.parametrize("stage", ["provider", "connect", "search", "tools/list"]) + async def test_task_cancellation_cleans_up(self, stage: str) -> None: + started = asyncio.Event() + original_connect = FakeTool.connect + + async def wait(*_args: Any, **_kwargs: Any) -> None: + started.set() + await asyncio.Event().wait() + + async def provider(_invocation: MCPToolInvocation) -> None: + if stage == "provider": + await wait() + + async def connect(tool: FakeTool) -> None: + await original_connect(tool) + if stage == "connect": + await wait() + + async def invoke_and_check_context() -> None: + try: + await handler.invoke_tool( + _invocation(tool_name="tools/list" if stage == "tools/list" else "search", headers={"X-Test": "1"}) + ) + finally: + assert handler._invocation_context.get() == () + + with ( + _patch_tool(), + patch.object(FakeTool, "connect", connect), + patch.object(FakeTool, "call_tool", new_callable=AsyncMock, side_effect=wait), + patch.object(FakeMcpSession, "list_tools", new_callable=AsyncMock, side_effect=wait), + ): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + task = asyncio.create_task(invoke_and_check_context()) + await asyncio.wait_for(started.wait(), timeout=5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not handler._active_invocations + assert len(FakeTool.instances) == (0 if stage == "provider" else 1) + for tool in FakeTool.instances: + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + + async def test_repeated_cancellation_waits_for_cleanup(self) -> None: + cleanup_started = asyncio.Event() + release_cleanup = asyncio.Event() + original_close = FakeTool.close + + async def close(tool: FakeTool) -> None: + cleanup_started.set() + await release_cleanup.wait() + await original_close(tool) + + with _patch_tool(), patch.object(FakeTool, "close", close): + async with DefaultMCPToolHandler(client_provider=AsyncMock(return_value=None)) as handler: + task = asyncio.create_task(handler.invoke_tool(_invocation(headers={"X-Test": "1"}))) + await asyncio.wait_for(cleanup_started.wait(), timeout=5) + for _ in range(2): + task.cancel() + await asyncio.sleep(0) + assert not task.done() + assert handler._active_invocations + release_cleanup.set() + with pytest.raises(asyncio.CancelledError): + await task + assert not handler._active_invocations + tool = FakeTool.instances[0] + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + + @pytest.mark.parametrize("stage", ["provider", "connect", "search", "tools/list", "close"]) + @pytest.mark.parametrize("cancel_shutdown", [False, True]) + async def test_shutdown_waits_for_invocation_cleanup(self, stage: str, cancel_shutdown: bool) -> None: + started = asyncio.Event() + release = asyncio.Event() + original_connect = FakeTool.connect + original_close = FakeTool.close + original_call = FakeTool.call_tool + original_list = FakeMcpSession.list_tools + + async def gate(phase: str) -> None: + if stage == phase: + started.set() + await release.wait() + + async def provider(_invocation: MCPToolInvocation) -> None: + await gate("provider") + + async def connect(tool: FakeTool) -> None: + await original_connect(tool) + await gate("connect") + + async def close(tool: FakeTool) -> None: + await gate("close") + await original_close(tool) + + async def call(tool: FakeTool, tool_name: str, **arguments: Any) -> Any: + await gate("search") + return await original_call(tool, tool_name, **arguments) + + async def list_tools(session: FakeMcpSession, params: Any = None) -> FakeListToolsResult: + await gate("tools/list") + return await original_list(session, params) + + provider_mock = AsyncMock(side_effect=provider) + with ( + _patch_tool(), + patch.object(FakeTool, "connect", connect), + patch.object(FakeTool, "close", close), + patch.object(FakeTool, "call_tool", call), + patch.object(FakeMcpSession, "list_tools", list_tools), + ): + handler = DefaultMCPToolHandler(client_provider=provider_mock) + invocation = _invocation( + tool_name="tools/list" if stage == "tools/list" else "search", headers={"X-Test": "1"} + ) + task = asyncio.create_task(handler.invoke_tool(invocation)) + await asyncio.wait_for(started.wait(), timeout=5) + shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert handler._closed + assert not shutdown.done() + rejected = await handler.invoke_tool(invocation) + assert rejected.is_error + assert "closed" in (rejected.error_message or "") + assert provider_mock.await_count == 1 + if cancel_shutdown: + shutdown.cancel() + with pytest.raises(asyncio.CancelledError): + await shutdown + second_shutdown = asyncio.create_task(handler.aclose()) + await asyncio.sleep(0) + assert not second_shutdown.done() + release.set() + result = await task + await second_shutdown + if not cancel_shutdown: + await shutdown + await handler.aclose() + assert result.is_error == (stage in ("provider", "connect")) + if result.is_error: + assert "closed" in (result.error_message or "") + assert not handler._active_invocations + assert not handler._cache + assert not handler._inflight + assert len(FakeTool.instances) == (0 if stage == "provider" else 1) + for tool in FakeTool.instances: + assert tool.close_count == 1 + assert tool._httpx_client is not None + assert tool._httpx_client.is_closed + + async def test_invalid_list_arguments_skip_provider(self) -> None: + provider = AsyncMock(return_value=None) + with _patch_tool(): + async with DefaultMCPToolHandler(client_provider=provider) as handler: + result = await handler.invoke_tool(_invocation(tool_name="tools/list", arguments={"q": "test"})) + assert result.is_error + assert "does not accept tool arguments" in (result.error_message or "") + provider.assert_not_awaited() + assert not FakeTool.instances # ---------- Construction --------------------------------------------------- From dcd2c4d1e0b0fce6bcebe0f3d244351ba9002a80 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:21:42 +0000 Subject: [PATCH 09/26] Allow disabling function tools on ModeProvider (#8458) --- .../Harness/AgentMode/AgentModeProvider.cs | 104 +++++++++--- .../AgentMode/AgentModeProviderOptions.cs | 18 +++ .../PublicAPI/net10.0/PublicAPI.Unshipped.txt | 5 + .../PublicAPI/net472/PublicAPI.Unshipped.txt | 5 + .../PublicAPI/net8.0/PublicAPI.Unshipped.txt | 5 + .../PublicAPI/net9.0/PublicAPI.Unshipped.txt | 5 + .../netstandard2.0/PublicAPI.Unshipped.txt | 5 + .../AgentMode/AgentModeProviderTests.cs | 153 ++++++++++++++++++ 8 files changed, 281 insertions(+), 19 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProvider.cs b/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProvider.cs index ac12296f995..f01ae98764b 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProvider.cs @@ -27,14 +27,18 @@ namespace Microsoft.Agents.AI; /// (autonomous execution). /// /// -/// This provider exposes the following tools to the agent: +/// By default, this provider exposes the following tools to the agent: /// /// mode_set — Switch the agent's operating mode. /// mode_get — Retrieve the agent's current operating mode. /// +/// Set or +/// to omit the corresponding built-in tool while +/// retaining mode state and workflow instructions. /// /// -/// Public helper methods and allow external code +/// Public helper methods and +/// allow external code /// to programmatically read and change the mode. /// /// @@ -44,16 +48,19 @@ namespace Microsoft.Agents.AI; /// public sealed class AgentModeProvider : AIContextProvider, IDisposable { + private const string ModeGetInstructions = "Use the mode_get tool to check your current operating mode.\n"; + private const string ModeSetInstructions = + "Use the mode_set tool to switch between modes as your work progresses. Only use mode_set if the user explicitly instructs/allows you to change modes.\n\n"; + private const string PlanModeTransition = + "7. When approval is granted, always switch to execute mode (using the `mode_set` tool), and follow the steps for *Execute mode*."; + private const string DefaultInstructions = """ ## Agent Mode - You can operate in different modes. Depending on the mode you are in, you will be required to follow different processes. - Use the mode_get tool to check your current operating mode. - Use the mode_set tool to switch between modes as your work progresses. Only use mode_set if the user explicitly instructs/allows you to change modes. - - You are currently operating in the {current_mode} mode. + {mode_get_instructions}{mode_set_instructions}You are currently operating in the {current_mode} mode. ### Mandatory Mode based Workflow @@ -80,7 +87,7 @@ 3. Do not proceed until you have received all the needed clarifications. 4. Do short exploratory research if it helps with being able to ask sensible clarifications from the user. 5. Write the plan to a memory file, so that it is retained even if compaction happens. Make sure to update the plan file if the user requests changes. 6. Present the plan to the user and ask for approval to switch to execute mode and process the plan. - 7. When approval is granted, always switch to execute mode (using the `mode_set` tool), and follow the steps for *Execute mode*. + {plan_mode_transition} """), new( "execute", @@ -102,7 +109,11 @@ 4. Mark tasks as completed as you finish them. private readonly ProviderSessionState _sessionState; private readonly IReadOnlyList _modes; private readonly string _defaultMode; - private readonly string? _instructions; + private readonly string _instructions; + private readonly bool _disableModeSetTool; + private readonly bool _disableModeGetTool; + private readonly bool _usesDefaultInstructions; + private readonly bool _usesDefaultModes; private readonly HashSet _validModeNames; private readonly string _modeNamesDisplay; private readonly ConditionalWeakTable _sessionLocks = new(); @@ -115,6 +126,7 @@ 4. Mark tasks as completed as you finish them. /// Optional settings that control provider behavior. When , defaults are used. public AgentModeProvider(AgentModeProviderOptions? options = null) { + this._usesDefaultModes = options?.Modes is null; this._modes = options?.Modes ?? s_defaultModes; if (this._modes.Count == 0) @@ -122,7 +134,10 @@ public AgentModeProvider(AgentModeProviderOptions? options = null) throw new ArgumentException("At least one mode must be configured.", nameof(options)); } + this._usesDefaultInstructions = options?.Instructions is null; this._instructions = options?.Instructions ?? DefaultInstructions; + this._disableModeSetTool = options?.DisableModeSetTool ?? false; + this._disableModeGetTool = options?.DisableModeGetTool ?? false; this._validModeNames = new HashSet(StringComparer.Ordinal); var modeNamesList = new List(this._modes.Count); @@ -202,7 +217,30 @@ public async Task GetModeAsync(AgentSession session, CancellationToken c /// A that represents the asynchronous operation. /// is . /// is not a configured mode. - public async Task SetModeAsync(AgentSession session, string mode, CancellationToken cancellationToken = default) + public Task SetModeAsync(AgentSession session, string mode, CancellationToken cancellationToken = default) => + this.SetModeAsync(session, mode, disableNotification: false, cancellationToken); + + /// + /// Sets the operating mode in the session state, optionally suppressing the mode-change notification. + /// + /// The agent session to update the mode in. + /// The new mode to set. + /// + /// to avoid notifying the agent about the mode change on its next invocation; + /// otherwise, . Use when the agent changes mode through + /// a custom function tool and has already observed the tool result. This also clears any pending + /// mode-change notification. + /// + /// The to monitor for cancellation requests. + /// A that represents the asynchronous operation. + /// is . + /// is not a configured mode. +#pragma warning disable RS0026 // The required Boolean parameter distinguishes this overload from the existing optional-cancellation overload. + public async Task SetModeAsync( + AgentSession session, + string mode, + bool disableNotification, + CancellationToken cancellationToken = default) { _ = Throw.IfNull(session); @@ -216,7 +254,11 @@ public async Task SetModeAsync(AgentSession session, string mode, CancellationTo string previousMode = state.CurrentMode; state.CurrentMode = mode; - if (!string.Equals(previousMode, mode, StringComparison.Ordinal)) + if (disableNotification) + { + state.PreviousModeForNotification = null; + } + else if (!string.Equals(previousMode, mode, StringComparison.Ordinal)) { state.PreviousModeForNotification = previousMode; } @@ -228,6 +270,7 @@ public async Task SetModeAsync(AgentSession session, string mode, CancellationTo sessionLock.Release(); } } +#pragma warning restore RS0026 /// protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) @@ -284,15 +327,32 @@ private string BuildInstructions(string currentMode) var modesListBuilder = new StringBuilder(); foreach (var mode in this._modes) { + string modeInstructions = mode.Instructions; + if (this._usesDefaultModes && mode.Name == "plan") + { + modeInstructions = modeInstructions.Replace( + "{plan_mode_transition}", + this._disableModeSetTool ? string.Empty : PlanModeTransition); + } + modesListBuilder.AppendLine($"#### {mode.Name}"); modesListBuilder.AppendLine(); - modesListBuilder.AppendLine(mode.Instructions.TrimEnd()); + modesListBuilder.AppendLine(modeInstructions.TrimEnd()); modesListBuilder.AppendLine(); } var modesListText = modesListBuilder.ToString(); + string instructions = this._instructions; + if (this._usesDefaultInstructions) + { + instructions = instructions + .Replace("{mode_get_instructions}", this._disableModeGetTool ? string.Empty : ModeGetInstructions) + .Replace( + "{mode_set_instructions}", + this._disableModeSetTool ? string.Empty : ModeSetInstructions); + } - return new StringBuilder(this._instructions) + return new StringBuilder(instructions) .Replace("{available_modes}", modesListText) .Replace("{current_mode}", currentMode) .ToString(); @@ -322,10 +382,11 @@ private SemaphoreSlim GetSessionLock(AgentSession? session) private AITool[] CreateTools(AgentSession? session) { var serializerOptions = AgentJsonUtilities.DefaultOptions; + var tools = new List(2); - return - [ - AIFunctionFactory.Create( + if (!this._disableModeSetTool) + { + tools.Add(AIFunctionFactory.Create( async (string mode) => { this.ValidateMode(mode); @@ -350,9 +411,12 @@ private AITool[] CreateTools(AgentSession? session) Name = "mode_set", Description = $"Switch the agent's operating mode. Supported modes: \"{this._modeNamesDisplay}\".", SerializerOptions = serializerOptions, - }), + })); + } - AIFunctionFactory.Create( + if (!this._disableModeGetTool) + { + tools.Add(AIFunctionFactory.Create( async () => { SemaphoreSlim sessionLock = this.GetSessionLock(session); @@ -371,7 +435,9 @@ private AITool[] CreateTools(AgentSession? session) Name = "mode_get", Description = "Get the agent's current operating mode.", SerializerOptions = serializerOptions, - }), - ]; + })); + } + + return tools.ToArray(); } } diff --git a/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProviderOptions.cs index 45471fba798..c711de29306 100644 --- a/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProviderOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Harness/AgentMode/AgentModeProviderOptions.cs @@ -42,6 +42,24 @@ public sealed class AgentModeProviderOptions /// public string? DefaultMode { get; set; } + /// + /// Gets or sets a value indicating whether the built-in mode_set tool is disabled. + /// + /// + /// When (the default), the provider exposes the mode_set tool. + /// When , the tool is not exposed, while mode state and instructions remain enabled. + /// + public bool DisableModeSetTool { get; set; } + + /// + /// Gets or sets a value indicating whether the built-in mode_get tool is disabled. + /// + /// + /// When (the default), the provider exposes the mode_get tool. + /// When , the tool is not exposed, while mode state and instructions remain enabled. + /// + public bool DisableModeGetTool { get; set; } + /// /// Represents an agent operating mode with a name and instructions. /// diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt index 26f303a9aa5..1ec4e243d09 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net10.0/PublicAPI.Unshipped.txt @@ -8,6 +8,11 @@ *REMOVED*[MAAI001]abstract Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.get -> System.TimeSpan [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.set -> void +Microsoft.Agents.AI.AgentModeProvider.SetModeAsync(Microsoft.Agents.AI.AgentSession! session, string! mode, bool disableNotification, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.set -> void +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.get -> string? [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.set -> void [MAAI001]const Microsoft.Agents.AI.FileAccessProvider.ReadLinesToolName = "file_access_read_lines" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt index 26f303a9aa5..1ec4e243d09 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net472/PublicAPI.Unshipped.txt @@ -8,6 +8,11 @@ *REMOVED*[MAAI001]abstract Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.get -> System.TimeSpan [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.set -> void +Microsoft.Agents.AI.AgentModeProvider.SetModeAsync(Microsoft.Agents.AI.AgentSession! session, string! mode, bool disableNotification, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.set -> void +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.get -> string? [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.set -> void [MAAI001]const Microsoft.Agents.AI.FileAccessProvider.ReadLinesToolName = "file_access_read_lines" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt index 26f303a9aa5..1ec4e243d09 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net8.0/PublicAPI.Unshipped.txt @@ -8,6 +8,11 @@ *REMOVED*[MAAI001]abstract Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.get -> System.TimeSpan [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.set -> void +Microsoft.Agents.AI.AgentModeProvider.SetModeAsync(Microsoft.Agents.AI.AgentSession! session, string! mode, bool disableNotification, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.set -> void +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.get -> string? [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.set -> void [MAAI001]const Microsoft.Agents.AI.FileAccessProvider.ReadLinesToolName = "file_access_read_lines" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt index 26f303a9aa5..1ec4e243d09 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/net9.0/PublicAPI.Unshipped.txt @@ -8,6 +8,11 @@ *REMOVED*[MAAI001]abstract Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.get -> System.TimeSpan [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.set -> void +Microsoft.Agents.AI.AgentModeProvider.SetModeAsync(Microsoft.Agents.AI.AgentSession! session, string! mode, bool disableNotification, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.set -> void +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.get -> string? [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.set -> void [MAAI001]const Microsoft.Agents.AI.FileAccessProvider.ReadLinesToolName = "file_access_read_lines" -> string! diff --git a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt index 26f303a9aa5..1ec4e243d09 100644 --- a/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt +++ b/dotnet/src/Microsoft.Agents.AI/PublicAPI/netstandard2.0/PublicAPI.Unshipped.txt @@ -8,6 +8,11 @@ *REMOVED*[MAAI001]abstract Microsoft.Agents.AI.AgentFileStore.SearchAsync(string! directory, string! regexPattern, string? globPattern = null, bool recursive = false, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task!>! [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.get -> System.TimeSpan [MAAI001]Microsoft.Agents.AI.BackgroundAgentsProviderOptions.WaitTimeout.set -> void +Microsoft.Agents.AI.AgentModeProvider.SetModeAsync(Microsoft.Agents.AI.AgentSession! session, string! mode, bool disableNotification, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) -> System.Threading.Tasks.Task! +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeGetTool.set -> void +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.get -> bool +Microsoft.Agents.AI.AgentModeProviderOptions.DisableModeSetTool.set -> void [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.get -> string? [MAAI001]Microsoft.Agents.AI.FileLineEdit.ExpectedLine.set -> void [MAAI001]const Microsoft.Agents.AI.FileAccessProvider.ReadLinesToolName = "file_access_read_lines" -> string! diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/AgentMode/AgentModeProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/AgentMode/AgentModeProviderTests.cs index cb38565780b..bbc1c2a7c99 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/AgentMode/AgentModeProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Harness/AgentMode/AgentModeProviderTests.cs @@ -40,6 +40,47 @@ public async Task ProvideAIContextAsync_ReturnsToolsAndInstructionsAsync() Assert.Equal(2, result.Tools!.Count()); } + /// + /// Verify that each built-in mode tool can be disabled independently. + /// + [Theory] + [InlineData(false, false, "mode_set,mode_get")] + [InlineData(false, true, "mode_set")] + [InlineData(true, false, "mode_get")] + [InlineData(true, true, "")] + public async Task ProvideAIContextAsync_DisablesConfiguredToolsAsync( + bool disableModeSetTool, + bool disableModeGetTool, + string expectedToolNames) + { + // Arrange + var provider = new AgentModeProvider(new AgentModeProviderOptions + { + DisableModeSetTool = disableModeSetTool, + DisableModeGetTool = disableModeGetTool, + }); + var agent = new Mock().Object; + var session = new ChatClientAgentSession(); +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(agent, session, new AIContext()); +#pragma warning restore MAAI001 + + // Act + AIContext result = await provider.InvokingAsync(context); + + // Assert + string[] expectedNames = expectedToolNames.Length == 0 ? [] : expectedToolNames.Split(','); + Assert.Equal(expectedNames, result.Tools!.Cast().Select(tool => tool.Name)); + Assert.Equal(!disableModeSetTool, result.Instructions!.Contains("mode_set", StringComparison.Ordinal)); + Assert.Equal(!disableModeGetTool, result.Instructions.Contains("mode_get", StringComparison.Ordinal)); + Assert.Contains("### Mandatory Mode based Workflow", result.Instructions); + Assert.Contains("You are currently operating in the plan mode.", result.Instructions); + Assert.DoesNotContain("{mode_get_instructions}", result.Instructions); + Assert.DoesNotContain("{mode_set_instructions}", result.Instructions); + Assert.DoesNotContain("{plan_mode_transition}", result.Instructions); + Assert.Equal("plan", await provider.GetModeAsync(session)); + } + /// /// Verify that the instructions include the current mode. /// @@ -198,6 +239,53 @@ public async Task PublicSetMode_ChangesModeAsync() Assert.Equal("execute", mode); } + /// + /// Verify that the public SetMode helper can suppress the mode-change notification. + /// + [Fact] + public async Task PublicSetMode_DisableNotification_DoesNotInjectNotificationAsync() + { + // Arrange + var provider = new AgentModeProvider(); + var agent = new Mock().Object; + var session = new ChatClientAgentSession(); +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(agent, session, new AIContext()); +#pragma warning restore MAAI001 + + // Act + await provider.SetModeAsync(session, "execute", disableNotification: true); + AIContext result = await provider.InvokingAsync(context); + + // Assert + Assert.Equal("execute", await provider.GetModeAsync(session)); + Assert.Null(result.Messages); + } + + /// + /// Verify that suppressing a notification clears an earlier pending mode-change notification. + /// + [Fact] + public async Task PublicSetMode_DisableNotification_ClearsPendingNotificationAsync() + { + // Arrange + var provider = new AgentModeProvider(); + var agent = new Mock().Object; + var session = new ChatClientAgentSession(); +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(agent, session, new AIContext()); +#pragma warning restore MAAI001 + await provider.SetModeAsync(session, "execute"); + + // Act + await provider.SetModeAsync(session, "plan", disableNotification: true); + AIContext result = await provider.InvokingAsync(context); + + // Assert + Assert.Equal("plan", await provider.GetModeAsync(session)); + Assert.Null(result.Messages); + } + /// /// Verify that the public SetMode helper throws for an unsupported value and does not persist the mode. /// @@ -303,6 +391,39 @@ public async Task Options_CustomInstructions_OverridesDefaultAsync() Assert.Equal("Custom mode instructions.", result.Instructions); } + /// + /// Verify that disabling tools does not rewrite custom instructions or custom mode guidance. + /// + [Fact] + public async Task Options_DisabledTools_PreserveCustomInstructionsAsync() + { + // Arrange + var options = new AgentModeProviderOptions + { + DisableModeSetTool = true, + DisableModeGetTool = true, + Instructions = "Use custom mode_set and mode_get behavior in {current_mode}.\n{available_modes}", + Modes = + [ + new AgentModeProviderOptions.AgentMode("draft", "Custom mode_set and mode_get guidance."), + ], + }; + var provider = new AgentModeProvider(options); + var agent = new Mock().Object; + var session = new ChatClientAgentSession(); +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(agent, session, new AIContext()); +#pragma warning restore MAAI001 + + // Act + AIContext result = await provider.InvokingAsync(context); + + // Assert + Assert.Empty(result.Tools!); + Assert.Contains("Use custom mode_set and mode_get behavior in draft.", result.Instructions); + Assert.Contains("Custom mode_set and mode_get guidance.", result.Instructions); + } + /// /// Verify that custom modes are used. /// @@ -535,6 +656,38 @@ public async Task ExternalModeChange_InjectsNotificationMessageAsync() Assert.Contains("execute", message.Text); } + /// + /// Verify that disabling both tools preserves mode state, instructions, and external-change notifications. + /// + [Fact] + public async Task DisabledTools_PreserveStateInstructionsAndNotificationsAsync() + { + // Arrange + var provider = new AgentModeProvider(new AgentModeProviderOptions + { + DisableModeSetTool = true, + DisableModeGetTool = true, + }); + var agent = new Mock().Object; + var session = new ChatClientAgentSession(); +#pragma warning disable MAAI001 + var context = new AIContextProvider.InvokingContext(agent, session, new AIContext()); +#pragma warning restore MAAI001 + _ = await provider.InvokingAsync(context); + + // Act + await provider.SetModeAsync(session, "execute"); + AIContext result = await provider.InvokingAsync(context); + + // Assert + Assert.Empty(result.Tools!); + Assert.Contains("You are currently operating in the execute mode.", result.Instructions); + Assert.Equal("execute", await provider.GetModeAsync(session)); + ChatMessage message = Assert.Single(result.Messages!); + Assert.Contains("plan", message.Text); + Assert.Contains("execute", message.Text); + } + /// /// Verify that the notification is only injected once (cleared after first read). /// From 81c47e295d8eb72275e227e176d79b77cfbff361 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:30:28 +0000 Subject: [PATCH 10/26] Clarify Agent Skills caching behavior (#8474) Document the default shared cache behavior for custom skill sources and point callers to the existing caching controls. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../Skills/AgentSkillsProviderBuilder.cs | 24 +++++++++++++++++++ .../CachingAgentSkillsSourceOptions.cs | 6 +++++ 2 files changed, 30 insertions(+) diff --git a/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillsProviderBuilder.cs b/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillsProviderBuilder.cs index e8ccf9ffb81..c172af407a0 100644 --- a/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillsProviderBuilder.cs +++ b/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillsProviderBuilder.cs @@ -130,6 +130,14 @@ public AgentSkillsProviderBuilder UseSkills(IEnumerable skills) /// overload, which creates a fresh /// source per build, or pass the source directly to an constructor /// with ownsSource: false to retain ownership. + /// + /// Custom sources are cached by default. Unless + /// returns a non-null key, + /// the cached skill list is shared by all invocations of the provider. If this source returns + /// skills that vary by session, user, tenant, or another security boundary, call + /// or configure an appropriate cache isolation key through + /// . + /// /// /// The custom skill source. /// This builder instance for chaining. @@ -145,6 +153,14 @@ public AgentSkillsProviderBuilder UseSource(AgentSkillsSource source) /// at build time. Use this overload when the source needs logging and should not require the /// caller to pass an explicitly. /// + /// + /// Custom sources are cached by default. Unless + /// returns a non-null key, + /// the cached skill list is shared by all invocations of the provider. If the created source + /// returns skills that vary by session, user, tenant, or another security boundary, call + /// or configure an appropriate cache isolation key through + /// . + /// /// A factory that creates the skill source given an optional logger factory. /// This builder instance for chaining. public AgentSkillsProviderBuilder UseSource(Func factory) @@ -247,6 +263,14 @@ public AgentSkillsProviderBuilder UseCachingOptions(Action + /// The source pipeline is cached by default. Unless + /// returns a non-null key, + /// the cached skill list is shared by all invocations of the provider. Sources that return skills + /// that vary by session, user, tenant, or another security boundary must use + /// or configure an appropriate cache isolation key through + /// . + /// + /// /// Build more than one provider from the same builder only when every source it produces is /// independent per build (for example, sources added via /// ). A source captured as a shared diff --git a/dotnet/src/Microsoft.Agents.AI/Skills/Decorators/CachingAgentSkillsSourceOptions.cs b/dotnet/src/Microsoft.Agents.AI/Skills/Decorators/CachingAgentSkillsSourceOptions.cs index 81f50067150..f942de4def1 100644 --- a/dotnet/src/Microsoft.Agents.AI/Skills/Decorators/CachingAgentSkillsSourceOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Skills/Decorators/CachingAgentSkillsSourceOptions.cs @@ -19,6 +19,12 @@ public sealed class CachingAgentSkillsSourceOptions /// the skills are cached under that key. /// /// + /// The shared cache bucket is used across all invocations of the provider. If the inner source + /// returns skills that vary by session, user, tenant, or another security boundary, configure a + /// key that provides the required isolation or disable caching through + /// . + /// + /// /// The isolation key should be low-cardinality and stable. /// High-cardinality keys (for example, per-session IDs) can cause the cache to grow without bound. /// From 6817fbdbf70d721726830b2fd70011ad5261e465 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 17 Sep 2026 16:34:59 +0000 Subject: [PATCH 11/26] docs: clarifies the additional languages policy (#8477) --- CONTRIBUTING.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 57454509db3..554640c23c3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -4,6 +4,16 @@ You can contribute to Agent Framework with issues and pull requests (PRs). Simpl filing issues for problems you encounter is a great way to contribute. Contributing code is greatly appreciated. +This repository is dedicated to the canonical .NET and Python implementations of +Microsoft Agent Framework. Microsoft contributors outside the product team should +engage with the maintainers before submitting pull requests that add new language +implementations or before starting a new Microsoft-owned repository for another +language. + +Contributors outside Microsoft are welcome to start their own repositories to +implement Microsoft Agent Framework in other languages, as long as they make it +clear that the effort is not owned or maintained directly by Microsoft. + ## Reporting Issues We always welcome bug reports, API proposals and overall feedback. Here are a few From 3c8344163d57fedd30559c255c5c9f4b6bb93dd6 Mon Sep 17 00:00:00 2001 From: Vincent Biret Date: Thu, 17 Sep 2026 16:39:20 +0000 Subject: [PATCH 12/26] ci: excludes samples from dotnet dependabot because they are causing a timeout (#8476) * ci: excludes samples from dotnet dependabot because they are causing a timeout * ci: aligns path specification --- .github/dependabot.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 92e64d5d631..2664f71d40b 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -8,6 +8,12 @@ updates: # Maintain dependencies for nuget - package-ecosystem: "nuget" directory: "dotnet/" + # The samples tree contains most of the .NET package manifests and causes + # the NuGet updater to exceed Dependabot's fixed execution timeout. Keep the + # production source, tests, and central package versions in scope so updates + # can be opened again. + exclude-paths: + - "samples/**" schedule: interval: "cron" cronjob: "0 8 * * 4,0" # Every Thursday(4) and Sunday(0) at 8:00 UTC From d188e8e6334881ce591b84b255ae9fb344e4663e Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+rogerbarreto@users.noreply.github.com> Date: Thu, 17 Sep 2026 16:53:15 +0000 Subject: [PATCH 13/26] .NET: Make Foundry user identity session sticky (#8434) --- .../UserIsolationAgent/Program.cs | 10 +-- .../UserIsolationAgent/README.md | 10 ++- .../FoundryAgent.cs | 21 ++++- .../FoundryAgentSessionExtensions.cs | 40 ++++++++- .../FoundryChatOptionsExtensions.cs | 71 +-------------- .../FoundryHostedRequestAgent.cs | 12 +-- .../UserIdentityPolicy.cs | 4 +- .../UserIdentityScope.cs | 4 +- .../HostedSessionAndUserIdentityTests.cs | 49 ++++++----- .../README.md | 4 +- .../FoundryHostedRequestTests.cs | 87 +++++++++++-------- 11 files changed, 164 insertions(+), 148 deletions(-) diff --git a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/Program.cs b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/Program.cs index b26b02f01dd..7150e882f7c 100644 --- a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/Program.cs +++ b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/Program.cs @@ -11,7 +11,6 @@ using DotNetEnv; using Microsoft.Agents.AI; using Microsoft.Agents.AI.Foundry; -using Microsoft.Extensions.AI; Env.TraversePath().Load(); @@ -96,22 +95,19 @@ User Isolation Agent Sample if (!userSessions.TryGetValue(userId, out ChatClientAgentSession? userSession)) { userSession = await agent.CreateFoundryHostedAgentSessionAsync( - hostedSessionId: hostedSessionId); + hostedSessionId: hostedSessionId, + userIdentity: userId); userSessions.Add(userId, userSession); Console.WriteLine($"Created an independent conversation for '{userId}'."); } - var runOptions = new ChatClientAgentRunOptions( - new ChatOptions().WithFoundryHostedAgentUserIdentity(userId)); - Console.ForegroundColor = ConsoleColor.Yellow; Console.Write($"Agent for {userId}> "); Console.ResetColor(); await foreach (AgentResponseUpdate update in agent.RunStreamingAsync( input, - userSession, - runOptions)) + userSession)) { Console.Write(update); } diff --git a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/README.md b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/README.md index 8788c3ce7db..9a9abe52b97 100644 --- a/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/README.md +++ b/dotnet/samples/04-hosting/FoundryHostedAgents/responses/Using-Samples/UserIsolationAgent/README.md @@ -4,9 +4,9 @@ This client demonstrates delegated user identity for a middle tier that serves s users through one Foundry hosted session. The sample creates one shared hosted session. Each user receives a separate Agent Framework -`AgentSession`, and every invocation sends that user's stable identifier through -`ChatOptions.WithFoundryHostedAgentUserIdentity`. Agent Framework places the value in the -`x-ms-user-identity` request header. +`AgentSession` created with that user's stable identifier. Agent Framework stores the identifier +with the session and automatically places it in the `x-ms-user-identity` request header on every +invocation that reuses the session. ## What Foundry isolates @@ -25,6 +25,10 @@ The sample deliberately creates one `AgentSession` per user. Reusing one `AgentS users would also reuse its conversation continuation identifier. Foundry rejects that cross user continuation rather than exposing the first user's history. +The delegated identity is fixed when `CreateFoundryHostedAgentSessionAsync` creates the local +`AgentSession`. To serve another user, create another `AgentSession`; it may still reference the same +hosted sandbox. + ## What the application must isolate Foundry isolates the conversation history it manages. It does not automatically partition arbitrary diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgent.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgent.cs index cee29456f44..893c68b77c2 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgent.cs @@ -164,7 +164,7 @@ public ValueTask CreateSessionAsync(string conversationId, Cancell /// /// Creates a local optionally pinned to a Foundry hosted-agent - /// session id (sandbox) and/or a server conversation id. + /// session id (sandbox), a server conversation id, and a delegated application user identity. /// /// /// Optional existing hosted-agent session id to pin on the session. The id identifies a Foundry @@ -181,8 +181,14 @@ public ValueTask CreateSessionAsync(string conversationId, Cancell /// history and hosted-agent session (sandbox) are separate Foundry concepts; see /// Sessions and conversations. /// + /// + /// Optional opaque application user identifier to associate with the session. When set, it is + /// stored in and sent as x-ms-user-identity on every + /// run that reuses this session. Create a separate for each user; + /// separate sessions may share the same . + /// /// The to monitor for cancellation requests. - /// A with the optional pins applied. + /// A with the optional bindings applied. /// /// /// The hosted-agent session itself is owned and lifecycle managed by Foundry Agent Service @@ -199,6 +205,7 @@ public ValueTask CreateSessionAsync(string conversationId, Cancell public async Task CreateFoundryHostedAgentSessionAsync( string? hostedSessionId = null, string? conversationId = null, + string? userIdentity = null, CancellationToken cancellationToken = default) { AgentSession session = conversationId is null @@ -212,6 +219,12 @@ public async Task CreateFoundryHostedAgentSessionAsync( typed.FoundryHostedAgentSessionId = hostedSessionId; } + if (userIdentity is not null) + { + // Non-null values are treated as an explicit identity binding; whitespace is rejected by Set. + typed.FoundryHostedAgentUserIdentity = userIdentity; + } + return typed; } @@ -297,11 +310,11 @@ private static AIAgent CreateResponsesChatClientAgent( } /// - /// Registers Foundry per-call pipeline policies and wraps the agent so request-scoped + /// Registers Foundry pipeline policies and wraps the agent so request context /// headers/body fields reach the wire: /// /// x-client-* via / - /// x-ms-user-identity and sticky agent_session_id via + /// sticky x-ms-user-identity and agent_session_id via /// /// Idempotent per decorator type. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgentSessionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgentSessionExtensions.cs index 7338b73c186..3383001f484 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgentSessionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryAgentSessionExtensions.cs @@ -18,6 +18,10 @@ namespace Microsoft.Agents.AI; /// serializing with the session. /// /// +/// The delegated user identity is stored in the same state bag. A session created for one user must +/// not be reused for another user, even when both sessions share the same hosted-agent session id. +/// +/// /// This is not . Per-call /// overrides use /// . @@ -31,6 +35,8 @@ public static class FoundryAgentSessionExtensions /// public const string FoundryHostedAgentSessionIdKey = "Microsoft.Agents.AI.Foundry.HostedAgentSessionId"; + private const string FoundryHostedAgentUserIdentityKey = "Microsoft.Agents.AI.Foundry.UserIdentity"; + extension(AgentSession session) { /// @@ -50,7 +56,7 @@ public static class FoundryAgentSessionExtensions /// /// /// Prefer creating or pinning through - /// . + /// . /// The property is populated automatically when Foundry creates a sandbox on first use. /// See /// Manage hosted agent sessions. @@ -73,5 +79,37 @@ internal set session.StateBag.SetValue(FoundryHostedAgentSessionIdKey, value); } } + + /// + /// Gets the delegated application user identity associated with this Agent Framework session. + /// + /// + /// The opaque application user identifier sent as x-ms-user-identity, or + /// when the session has no delegated identity. + /// + /// + /// The identity is fixed when the session is created through + /// . + /// Reusing this session automatically sends the same identity on every run. Create a separate + /// for each user; separate sessions may share the same + /// FoundryHostedAgentSessionId. + /// + public string? FoundryHostedAgentUserIdentity + { + get + { + _ = Throw.IfNull(session); + return session.StateBag.TryGetValue(FoundryHostedAgentUserIdentityKey, out var value) + ? value + : null; + } + + internal set + { + _ = Throw.IfNull(session); + _ = Throw.IfNullOrWhitespace(value); + session.StateBag.SetValue(FoundryHostedAgentUserIdentityKey, value); + } + } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryChatOptionsExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryChatOptionsExtensions.cs index 8f6f00674e7..b91e3fbba26 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryChatOptionsExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryChatOptionsExtensions.cs @@ -13,18 +13,15 @@ namespace Microsoft.Extensions.AI; /// /// /// -/// Use these helpers to attach per-call Foundry request fields: -/// -/// sends agent_session_id on the Responses body. -/// sends x-ms-user-identity on the request. -/// +/// Use to send agent_session_id on the +/// Responses body for a single run. /// /// /// Hosted-agent session ids supplied via participate in the same /// conflict rule as : if the already /// holds a different hosted id in its , the run throws /// . Prefer pinning at session creation via -/// . +/// . /// /// [Experimental(DiagnosticIds.Experiments.AIOpenAIRequestPolicies)] @@ -39,12 +36,6 @@ public static class FoundryChatOptionsExtensions /// internal const string FoundryHostedAgentSessionIdKey = "Microsoft.Agents.AI.Foundry.HostedAgentSessionId"; - /// - /// Well-known key used to carry the per-call - /// user identity value. - /// - internal const string FoundryHostedAgentUserIdentityKey = "Microsoft.Agents.AI.Foundry.UserIdentity"; - /// /// Attaches a hosted-agent session id to the per-call carrier. /// @@ -52,7 +43,7 @@ public static class FoundryChatOptionsExtensions /// /// Only valid when the run's session has no hosted id yet, or already has this same id. /// Prefer - /// + /// /// to pin at session creation. /// /// @@ -71,44 +62,6 @@ public static ChatOptions WithFoundryHostedAgentSessionId(this ChatOptions optio return options; } - /// - /// Attaches a delegated user identity value that will be sent as the - /// x-ms-user-identity request header. - /// - /// The per-call chat options to mutate. - /// Opaque application user identifier. Must be non-empty. - /// for fluent chaining. - /// - /// - /// User identity is always request-scoped. It is never stored on . - /// - /// - /// Per Foundry hosted-agent isolation, a Responses chain created under one user cannot be - /// continued by another user via previous_response_id, even when both calls share the - /// same hosted sandbox (agent_session_id). See - /// Multiplex multiple users in one hosted agent session. - /// Reusing one across identities typically reuses that chain, so the - /// second identity's run fails at the platform (observed as a response not-found error). Prefer - /// a distinct per identity; those sessions may still share one hosted - /// sandbox pin via or - /// . - /// - /// - /// The value is stored in . Replacing that - /// dictionary after calling this method removes the value; populate or replace the dictionary - /// first, then call this method. - /// - /// - public static ChatOptions WithFoundryHostedAgentUserIdentity(this ChatOptions options, string userIdentity) - { - _ = Throw.IfNull(options); - _ = Throw.IfNullOrWhitespace(userIdentity); - - options.AdditionalProperties ??= new AdditionalPropertiesDictionary(); - options.AdditionalProperties[FoundryHostedAgentUserIdentityKey] = userIdentity; - return options; - } - /// Reads the per-call hosted-agent session id stamped by . internal static string? GetFoundryHostedAgentSessionId(this ChatOptions options) { @@ -124,20 +77,4 @@ public static ChatOptions WithFoundryHostedAgentUserIdentity(this ChatOptions op return raw as string; } - - /// Reads the per-call user identity stamped by . - internal static string? GetFoundryHostedAgentUserIdentity(this ChatOptions options) - { - if (options.AdditionalProperties is null) - { - return null; - } - - if (!options.AdditionalProperties.TryGetValue(FoundryHostedAgentUserIdentityKey, out var raw)) - { - return null; - } - - return raw as string; - } } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryHostedRequestAgent.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryHostedRequestAgent.cs index 7b4dabd4f1e..576cb770138 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryHostedRequestAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/FoundryHostedRequestAgent.cs @@ -15,10 +15,10 @@ namespace Microsoft.Agents.AI.Foundry; /// -/// Delegating agent that applies Foundry hosted-agent request context per run: +/// Delegating agent that applies Foundry hosted-agent request context: /// resolves the sticky hosted-agent session id, injects agent_session_id into the -/// Responses body, stamps x-ms-user-identity, and writes the platform-returned session -/// id back onto the . +/// Responses body, stamps the session's x-ms-user-identity, and writes the +/// platform-returned session id back onto the . /// internal sealed class FoundryHostedRequestAgent : DelegatingAIAgent { @@ -95,9 +95,9 @@ The hosted-agent session id provided via ChatOptions is different from the id st var effectiveOptions = EnsureChatOptions(options, out chatOptions); AttachHostedSessionIdFactory(chatOptions, sessionIdBox); - // Always assign (including null) so a nested Foundry run that omits the per-call Foundry - // user identity does not inherit a parent AsyncLocal value and stamp the wrong header. - UserIdentityScope.Current = chatOptions.GetFoundryHostedAgentUserIdentity(); + // Always assign (including null) so a nested Foundry run without a session identity does + // not inherit a parent AsyncLocal value and stamp the wrong header. + UserIdentityScope.Current = session?.FoundryHostedAgentUserIdentity; return new PreparedRun(effectiveOptions, sessionIdBox); } diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityPolicy.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityPolicy.cs index a71655e2bc1..4fa86673df4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityPolicy.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityPolicy.cs @@ -7,8 +7,8 @@ namespace Microsoft.Agents.AI.Foundry; /// -/// Pipeline policy that stamps x-ms-user-identity from -/// onto outbound OpenAI Responses requests. +/// Pipeline policy that stamps the current session's x-ms-user-identity from +/// onto outbound OpenAI Responses requests. /// internal sealed class UserIdentityPolicy : PipelinePolicy { diff --git a/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityScope.cs b/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityScope.cs index b7e5630eb7b..0e116251308 100644 --- a/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityScope.cs +++ b/dotnet/src/Microsoft.Agents.AI.Foundry/UserIdentityScope.cs @@ -5,14 +5,14 @@ namespace Microsoft.Agents.AI.Foundry; /// -/// AsyncLocal carrier for the per-call x-ms-user-identity value from +/// AsyncLocal carrier for the session's x-ms-user-identity value from /// to . /// internal static class UserIdentityScope { private static readonly AsyncLocal s_current = new(); - /// Gets or sets the per-async-flow user identity value. + /// Gets or sets the user identity value for the current asynchronous flow. public static string? Current { get => s_current.Value; diff --git a/dotnet/tests/Foundry.Hosting.IntegrationTests/HostedSessionAndUserIdentityTests.cs b/dotnet/tests/Foundry.Hosting.IntegrationTests/HostedSessionAndUserIdentityTests.cs index 3ea40e425ab..871e5442045 100644 --- a/dotnet/tests/Foundry.Hosting.IntegrationTests/HostedSessionAndUserIdentityTests.cs +++ b/dotnet/tests/Foundry.Hosting.IntegrationTests/HostedSessionAndUserIdentityTests.cs @@ -15,14 +15,13 @@ using Foundry.Hosting.IntegrationTests.Fixtures; using Microsoft.Agents.AI; using Microsoft.Agents.AI.Foundry; -using Microsoft.Extensions.AI; using Shared.IntegrationTests; namespace Foundry.Hosting.IntegrationTests; /// -/// Live tests for client-side Foundry hosted session sticky behavior and per-call -/// x-ms-user-identity pass-through against a real hosted agent. +/// Live tests for client-side Foundry hosted session and delegated user identity sticky +/// behavior against a real hosted agent. /// /// /// @@ -122,18 +121,21 @@ public async Task SameHostedSandbox_DifferentAgentSessionsAndUserIdentities_Yiel try { // Act: alice creates the sandbox via service-managed sticky capture. - ChatClientAgentSession aliceSession = await agent.CreateFoundryHostedAgentSessionAsync(); - string aliceUserId = await this.RunAndReadUserIdAsync(agent, aliceSession, "alice-it"); + ChatClientAgentSession aliceSession = await agent.CreateFoundryHostedAgentSessionAsync( + userIdentity: "alice-it"); + string aliceUserId = await this.RunAndReadUserIdAsync(agent, aliceSession); hostedSessionId = aliceSession.FoundryHostedAgentSessionId; Assert.False(string.IsNullOrWhiteSpace(hostedSessionId)); string? aliceConversationId = aliceSession.ConversationId; // Act: bob gets a fresh AgentSession pinned to the same hosted sandbox. - ChatClientAgentSession bobSession = await agent.CreateFoundryHostedAgentSessionAsync(hostedSessionId: hostedSessionId); + ChatClientAgentSession bobSession = await agent.CreateFoundryHostedAgentSessionAsync( + hostedSessionId: hostedSessionId, + userIdentity: "bob-it"); Assert.NotSame(aliceSession, bobSession); Assert.Equal(hostedSessionId, bobSession.FoundryHostedAgentSessionId); - string bobUserId = await this.RunAndReadUserIdAsync(agent, bobSession, "bob-it"); + string bobUserId = await this.RunAndReadUserIdAsync(agent, bobSession); // Assert: hosted sandbox stays the same on both sessions after bob's response. Assert.Equal(hostedSessionId, aliceSession.FoundryHostedAgentSessionId); @@ -154,7 +156,9 @@ aliceConversationId is not null // Assert: platform user keys differ for alice vs bob. Assert.NotEqual("missing", aliceUserId); Assert.NotEqual("missing", bobUserId); - Assert.NotEqual(aliceUserId, bobUserId); + Assert.False( + string.Equals(aliceUserId, bobUserId, StringComparison.Ordinal), + "Expected different platform user keys for separate delegated identities."); } finally { @@ -163,25 +167,34 @@ aliceConversationId is not null } [Fact(Skip = "Requires live Foundry hosted agent image, bootstrap it-user-identity, and delegation permission for x-ms-user-identity.")] - public async Task SameSession_SameUserIdentity_YieldsStablePlatformUserIdAsync() + public async Task SerializedSession_SameUserIdentity_YieldsStablePlatformUserIdAsync() { // Arrange FoundryAgent agent = this.CreateFoundryAgent(); - ChatClientAgentSession session = await agent.CreateFoundryHostedAgentSessionAsync(); + ChatClientAgentSession session = await agent.CreateFoundryHostedAgentSessionAsync( + userIdentity: "stable-user-it"); string? hostedSessionId = null; try { // Act - string first = await this.RunAndReadUserIdAsync(agent, session, "stable-user-it"); + string first = await this.RunAndReadUserIdAsync(agent, session); hostedSessionId = session.FoundryHostedAgentSessionId; - string second = await this.RunAndReadUserIdAsync(agent, session, "stable-user-it"); + + var serializedSession = await agent.SerializeSessionAsync(session); + session = Assert.IsType( + await agent.DeserializeSessionAsync(serializedSession)); + + string second = await this.RunAndReadUserIdAsync(agent, session); // Assert Assert.False(string.IsNullOrWhiteSpace(hostedSessionId)); Assert.Equal(hostedSessionId, session.FoundryHostedAgentSessionId); + Assert.Equal("stable-user-it", session.FoundryHostedAgentUserIdentity); Assert.NotEqual("missing", first); - Assert.Equal(first, second); + Assert.True( + string.Equals(first, second, StringComparison.Ordinal), + "Expected the restored session to keep the same platform user key."); } finally { @@ -189,19 +202,15 @@ public async Task SameSession_SameUserIdentity_YieldsStablePlatformUserIdAsync() } } - private async Task RunAndReadUserIdAsync(FoundryAgent agent, AgentSession session, string userIdentity) + private async Task RunAndReadUserIdAsync(FoundryAgent agent, AgentSession session) { - var options = new ChatClientAgentRunOptions( - new ChatOptions().WithFoundryHostedAgentUserIdentity(userIdentity)); - var response = await agent.RunAsync( "Acknowledge the request briefly.", - session, - options); + session); Assert.False(string.IsNullOrWhiteSpace(response.Text)); Match match = s_userIdToken.Match(response.Text); - Assert.True(match.Success, $"Expected USER-ID: token in response text. Actual: {response.Text}"); + Assert.True(match.Success, "Expected a USER-ID: token in the response text."); return match.Groups[1].Value; } diff --git a/dotnet/tests/Foundry.Hosting.IntegrationTests/README.md b/dotnet/tests/Foundry.Hosting.IntegrationTests/README.md index b29c8f85c46..1b5bedf84b2 100644 --- a/dotnet/tests/Foundry.Hosting.IntegrationTests/README.md +++ b/dotnet/tests/Foundry.Hosting.IntegrationTests/README.md @@ -43,8 +43,8 @@ supported hosting paths. - `CreateFoundryHostedAgentSessionAsync` sticky hosted `agent_session_id` (service-managed and admin `CreateSession` / `DeleteSession` pin) -- per-call `ChatOptions.WithFoundryHostedAgentUserIdentity` (`x-ms-user-identity`) producing distinct - platform user keys inside the container +- session-sticky delegated identity (`x-ms-user-identity`) producing distinct platform user keys + inside the container, including after session serialization and restoration The container scenario injects `USER-ID:` via `EchoPlatformUserIdContextProvider`, reading `HostedSessionContext.UserId` (from diff --git a/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/FoundryHostedRequestTests.cs b/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/FoundryHostedRequestTests.cs index 3dc2e738f43..aee7a5000be 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/FoundryHostedRequestTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Foundry.UnitTests/FoundryHostedRequestTests.cs @@ -19,7 +19,7 @@ namespace Microsoft.Agents.AI.Foundry.UnitTests; /// -/// Tests for hosted-agent session sticky behavior and per-call user identity. +/// Tests for hosted-agent session and delegated user identity sticky behavior. /// public sealed class FoundryHostedRequestTests { @@ -32,35 +32,55 @@ public void WithFoundryHostedAgentSessionId_WritesOptionsCarrier() } [Fact] - public void WithFoundryHostedAgentUserIdentity_WritesOptionsCarrier() - { - var options = new ChatOptions(); - options.WithFoundryHostedAgentUserIdentity("alice"); - Assert.Equal("alice", options.GetFoundryHostedAgentUserIdentity()); - } - - [Fact] - public async Task CreateFoundryHostedAgentSessionAsync_PinsHostedAndConversationIdsAsync() + public async Task CreateFoundryHostedAgentSessionAsync_PinsHostedConversationAndUserIdentityAsync() { FoundryAgent agent = CreateFoundryAgent(); ChatClientAgentSession session = await agent.CreateFoundryHostedAgentSessionAsync( hostedSessionId: "sess-1", - conversationId: "conv-1"); + conversationId: "conv-1", + userIdentity: "alice"); Assert.Equal("sess-1", session.FoundryHostedAgentSessionId); Assert.Equal("conv-1", session.ConversationId); + Assert.Equal("alice", session.FoundryHostedAgentUserIdentity); Assert.True(session.StateBag.TryGetValue(FoundryAgentSessionExtensions.FoundryHostedAgentSessionIdKey, out var raw)); Assert.Equal("sess-1", raw); } [Fact] - public async Task CreateFoundryHostedAgentSessionAsync_WithoutIds_LeavesBothEmptyAsync() + public async Task CreateFoundryHostedAgentSessionAsync_WithoutBindings_LeavesAllEmptyAsync() { FoundryAgent agent = CreateFoundryAgent(); ChatClientAgentSession session = await agent.CreateFoundryHostedAgentSessionAsync(); Assert.Null(session.FoundryHostedAgentSessionId); Assert.Null(session.ConversationId); + Assert.Null(session.FoundryHostedAgentUserIdentity); + } + + [Fact] + public async Task UserIdentity_RoundTripsWithSessionSerializationAsync() + { + FoundryAgent agent = CreateFoundryAgent(); + ChatClientAgentSession session = await agent.CreateFoundryHostedAgentSessionAsync( + hostedSessionId: "sess-1", + conversationId: "conv-1", + userIdentity: "alice"); + + JsonElement serialized = await agent.SerializeSessionAsync(session); + ChatClientAgentSession restored = Assert.IsType( + await agent.DeserializeSessionAsync(serialized)); + + Assert.Equal("sess-1", restored.FoundryHostedAgentSessionId); + Assert.Equal("conv-1", restored.ConversationId); + Assert.Equal("alice", restored.FoundryHostedAgentUserIdentity); + + string? seenIdentity = null; + var requestAgent = new FoundryHostedRequestAgent( + new ProbeAgent(onRun: _ => seenIdentity = UserIdentityScope.Current)); + await requestAgent.RunAsync("hi", restored); + + Assert.Equal("alice", seenIdentity); } [Fact] @@ -72,6 +92,15 @@ await Assert.ThrowsAsync( () => agent.CreateFoundryHostedAgentSessionAsync(hostedSessionId: " ")); } + [Fact] + public async Task CreateFoundryHostedAgentSessionAsync_WhitespaceUserIdentity_ThrowsAsync() + { + FoundryAgent agent = CreateFoundryAgent(); + + await Assert.ThrowsAsync( + () => agent.CreateFoundryHostedAgentSessionAsync(userIdentity: " ")); + } + [Fact] public async Task Conflict_SessionAndOptionsHostedIdsDiffer_ThrowsAsync() { @@ -136,27 +165,19 @@ public async Task OptionsHostedId_WhenSessionEmpty_IsInjectedAndStickyAfterRunAs } [Fact] - public async Task UserIdentity_DifferentPerCall_OnSameSession_IsAllowedAsync() + public async Task UserIdentity_SessionValueIsUsedAcrossRunsAsync() { - // Pipeline still allows different identities on one AgentSession (request-scoped header). - // On a live hosted agent, Foundry binds previous_response_id chains to the creating user, so - // prefer distinct AgentSessions per identity; sandbox id may still be shared. var seen = new List(); var inner = new ProbeAgent(onRun: _ => seen.Add(UserIdentityScope.Current)); var agent = new FoundryHostedRequestAgent(inner); var session = new TestSession(); session.FoundryHostedAgentSessionId = "sess-shared"; + session.FoundryHostedAgentUserIdentity = "alice"; - await agent.RunAsync( - "hi", - session, - new ChatClientAgentRunOptions(new ChatOptions().WithFoundryHostedAgentUserIdentity("alice"))); - await agent.RunAsync( - "hi", - session, - new ChatClientAgentRunOptions(new ChatOptions().WithFoundryHostedAgentUserIdentity("bob"))); + await agent.RunAsync("hi", session); + await agent.RunAsync("hi", session); - Assert.Equal(["alice", "bob"], seen); + Assert.Equal(["alice", "alice"], seen); Assert.Equal("sess-shared", session.FoundryHostedAgentSessionId); } @@ -166,13 +187,12 @@ public async Task UserIdentity_OmittedAfterParent_ClearsAsyncLocalScopeAsync() var seen = new List(); var inner = new ProbeAgent(onRun: _ => seen.Add(UserIdentityScope.Current)); var agent = new FoundryHostedRequestAgent(inner); - var session = new TestSession(); + var sessionWithIdentity = new TestSession(); + sessionWithIdentity.FoundryHostedAgentUserIdentity = "alice"; + var sessionWithoutIdentity = new TestSession(); - await agent.RunAsync( - "hi", - session, - new ChatClientAgentRunOptions(new ChatOptions().WithFoundryHostedAgentUserIdentity("alice"))); - await agent.RunAsync("hi", session, new ChatClientAgentRunOptions(new ChatOptions())); + await agent.RunAsync("hi", sessionWithIdentity); + await agent.RunAsync("hi", sessionWithoutIdentity); Assert.Equal(["alice", null], seen); } @@ -263,11 +283,10 @@ public async Task EndToEnd_UserIdentity_AndHostedSessionId_ReachWireAsync() AIAgent agent = new FoundryHostedRequestAgent(new ClientHeadersAgent(chatAgent)); AgentSession session = await chatAgent.CreateSessionAsync(); session.FoundryHostedAgentSessionId = "sess-pinned"; + session.FoundryHostedAgentUserIdentity = "alice"; var runOptions = new ChatClientAgentRunOptions( - new ChatOptions() - .WithFoundryHostedAgentUserIdentity("alice") - .WithClientHeader("x-client-end-user-id", "alice-app")); + new ChatOptions().WithClientHeader("x-client-end-user-id", "alice-app")); // Response returns a different hosted session id than the pin → unexpected switch. InvalidOperationException ex = await Assert.ThrowsAsync( From 8559a7b2980b25668b2a86a5e32bf572806a7e7d Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:03:32 +0000 Subject: [PATCH 14/26] Build(deps): Bump azure/login (#8485) Bumps [azure/login](https://github.com/azure/login) from 3.0.2 to 3.1.0. - [Release notes](https://github.com/azure/login/releases) - [Commits](https://github.com/azure/login/compare/7ddb5af1ef8758cf1353cf3b42f940aee27ba21c...a641126d1b8aa4d1fa005f4f92df94a3a4c4c906) --- updated-dependencies: - dependency-name: azure/login dependency-version: 3.1.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/actions/sample-validation-setup/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/sample-validation-setup/action.yml b/.github/actions/sample-validation-setup/action.yml index 020823448f0..73e0b69a144 100644 --- a/.github/actions/sample-validation-setup/action.yml +++ b/.github/actions/sample-validation-setup/action.yml @@ -59,7 +59,7 @@ runs: sample-playbooks-${{ github.job }}- - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ inputs.azure-client-id }} tenant-id: ${{ inputs.azure-tenant-id }} From 5e0f3a54f924e286f1ebdbac6bb520864290bb9c Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:04:37 +0000 Subject: [PATCH 15/26] Build(deps): Bump astral-sh/setup-uv from 10.0.1 to 10.1.0 (#8482) Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 10.0.1 to 10.1.0. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](https://github.com/astral-sh/setup-uv/compare/20cfd1bf945f4377ade1205e4dbc17946fc9a30d...bec219d24cd3e171d82865faccec33120bb574f4) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 10.1.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/devflow-pr-review.yml | 2 +- .github/workflows/github-automation-tests.yml | 2 +- .github/workflows/issue-triage.yml | 2 +- .github/workflows/python-api-compatibility.yml | 2 +- .github/workflows/python-docs.yml | 2 +- .github/workflows/python-lab-tests.yml | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.github/workflows/devflow-pr-review.yml b/.github/workflows/devflow-pr-review.yml index b15ba1482cc..4cdb2b6122e 100644 --- a/.github/workflows/devflow-pr-review.yml +++ b/.github/workflows/devflow-pr-review.yml @@ -246,7 +246,7 @@ jobs: python-version: "3.13" - name: Set up uv - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: version: "0.11.x" enable-cache: auto diff --git a/.github/workflows/github-automation-tests.yml b/.github/workflows/github-automation-tests.yml index a05f7bb8fcf..cc3fc9c52ec 100644 --- a/.github/workflows/github-automation-tests.yml +++ b/.github/workflows/github-automation-tests.yml @@ -41,7 +41,7 @@ jobs: with: python-version: "3.11" - - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + - uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: version: "0.12.9" diff --git a/.github/workflows/issue-triage.yml b/.github/workflows/issue-triage.yml index f06875c830c..fb044df84c3 100644 --- a/.github/workflows/issue-triage.yml +++ b/.github/workflows/issue-triage.yml @@ -190,7 +190,7 @@ jobs: python-version: "3.13" - name: Set up uv - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: version: "0.11.x" enable-cache: true diff --git a/.github/workflows/python-api-compatibility.yml b/.github/workflows/python-api-compatibility.yml index 117cb414a45..7e23827cfc0 100644 --- a/.github/workflows/python-api-compatibility.yml +++ b/.github/workflows/python-api-compatibility.yml @@ -31,7 +31,7 @@ jobs: persist-credentials: false - name: Set up uv - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: enable-cache: false version: "0.12.9" diff --git a/.github/workflows/python-docs.yml b/.github/workflows/python-docs.yml index 50fb14ad91b..0b4a4c4e638 100644 --- a/.github/workflows/python-docs.yml +++ b/.github/workflows/python-docs.yml @@ -26,7 +26,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up uv - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: version-file: "python/pyproject.toml" enable-cache: auto diff --git a/.github/workflows/python-lab-tests.yml b/.github/workflows/python-lab-tests.yml index 6cff1b23766..a5885092f0b 100644 --- a/.github/workflows/python-lab-tests.yml +++ b/.github/workflows/python-lab-tests.yml @@ -59,7 +59,7 @@ jobs: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Set up uv - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: version-file: "python/packages/lab/pyproject.toml" python-version: ${{ matrix.python-version }} From 200b78765d5a8eb78ffa661e331c4c83fceee736 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:05:04 +0000 Subject: [PATCH 16/26] Build(deps): Bump astral-sh/setup-uv in /.github/actions/python-setup (#8484) Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 10.0.1 to 10.1.0. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](https://github.com/astral-sh/setup-uv/compare/20cfd1bf945f4377ade1205e4dbc17946fc9a30d...bec219d24cd3e171d82865faccec33120bb574f4) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 10.1.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/actions/python-setup/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/python-setup/action.yml b/.github/actions/python-setup/action.yml index 2ceb4dd5cf9..02fd5b48042 100644 --- a/.github/actions/python-setup/action.yml +++ b/.github/actions/python-setup/action.yml @@ -17,7 +17,7 @@ runs: using: "composite" steps: - name: Set up uv - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 with: version-file: "python/pyproject.toml" enable-cache: true From 3b45c56e8894f5f719a8bf12f523bb4e1700ee82 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:05:34 +0000 Subject: [PATCH 17/26] Build(deps): Bump azure/login in /.github/actions/github-app-token (#8483) Bumps [azure/login](https://github.com/azure/login) from 3.0.2 to 3.1.0. - [Release notes](https://github.com/azure/login/releases) - [Commits](https://github.com/azure/login/compare/7ddb5af1ef8758cf1353cf3b42f940aee27ba21c...a641126d1b8aa4d1fa005f4f92df94a3a4c4c906) --- updated-dependencies: - dependency-name: azure/login dependency-version: 3.1.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/actions/github-app-token/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/github-app-token/action.yml b/.github/actions/github-app-token/action.yml index 519d9b44b5e..a5551f74588 100644 --- a/.github/actions/github-app-token/action.yml +++ b/.github/actions/github-app-token/action.yml @@ -59,7 +59,7 @@ runs: id: azure-login if: ${{ (inputs.mode || 'app-with-fallback') != 'pat' }} continue-on-error: true - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ inputs.azure-client-id }} tenant-id: ${{ inputs.azure-tenant-id }} From bc57ce8857a5a38af5d87068daf1677fa3a01d32 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:06:18 +0000 Subject: [PATCH 18/26] Build(deps): Bump azure/login from 3.0.2 to 3.1.0 (#8481) Bumps [azure/login](https://github.com/azure/login) from 3.0.2 to 3.1.0. - [Release notes](https://github.com/azure/login/releases) - [Commits](https://github.com/azure/login/compare/7ddb5af1ef8758cf1353cf3b42f940aee27ba21c...a641126d1b8aa4d1fa005f4f92df94a3a4c4c906) --- updated-dependencies: - dependency-name: azure/login dependency-version: 3.1.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/dotnet-build-and-test.yml | 4 ++-- .github/workflows/dotnet-integration-tests.yml | 2 +- .github/workflows/dotnet-verify-samples.yml | 2 +- .github/workflows/issue-triage.yml | 2 +- .github/workflows/python-integration-tests.yml | 6 +++--- .github/workflows/python-merge-tests.yml | 6 +++--- 6 files changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/workflows/dotnet-build-and-test.yml b/.github/workflows/dotnet-build-and-test.yml index 02ba53477d6..4345cd6fdaf 100644 --- a/.github/workflows/dotnet-build-and-test.yml +++ b/.github/workflows/dotnet-build-and-test.yml @@ -418,7 +418,7 @@ jobs: - name: Azure CLI Login if: github.event_name != 'pull_request' && matrix.integration-tests - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} @@ -562,7 +562,7 @@ jobs: run: dotnet build dotnet/tests/Foundry.Hosting.IntegrationTests/Foundry.Hosting.IntegrationTests.csproj -c "$configuration" --warnaserror - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} diff --git a/.github/workflows/dotnet-integration-tests.yml b/.github/workflows/dotnet-integration-tests.yml index 032f5dabcee..eef36745030 100644 --- a/.github/workflows/dotnet-integration-tests.yml +++ b/.github/workflows/dotnet-integration-tests.yml @@ -77,7 +77,7 @@ jobs: done - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} diff --git a/.github/workflows/dotnet-verify-samples.yml b/.github/workflows/dotnet-verify-samples.yml index 49865915ea8..3855914f6c9 100644 --- a/.github/workflows/dotnet-verify-samples.yml +++ b/.github/workflows/dotnet-verify-samples.yml @@ -58,7 +58,7 @@ jobs: - name: Azure CLI Login if: github.event_name != 'pull_request' - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} diff --git a/.github/workflows/issue-triage.yml b/.github/workflows/issue-triage.yml index fb044df84c3..cf0a2749585 100644 --- a/.github/workflows/issue-triage.yml +++ b/.github/workflows/issue-triage.yml @@ -200,7 +200,7 @@ jobs: run: uv sync --frozen - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} diff --git a/.github/workflows/python-integration-tests.yml b/.github/workflows/python-integration-tests.yml index 30c29ec14b1..d85bba7c9f9 100644 --- a/.github/workflows/python-integration-tests.yml +++ b/.github/workflows/python-integration-tests.yml @@ -233,7 +233,7 @@ jobs: python-version: ${{ env.UV_PYTHON }} os: ${{ runner.os }} - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} @@ -403,7 +403,7 @@ jobs: python-version: ${{ env.UV_PYTHON }} os: ${{ runner.os }} - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} @@ -453,7 +453,7 @@ jobs: python-version: ${{ env.UV_PYTHON }} os: ${{ runner.os }} - name: Azure CLI Login - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} diff --git a/.github/workflows/python-merge-tests.yml b/.github/workflows/python-merge-tests.yml index 1891b4dcc0e..2d75ad02844 100644 --- a/.github/workflows/python-merge-tests.yml +++ b/.github/workflows/python-merge-tests.yml @@ -348,7 +348,7 @@ jobs: os: ${{ runner.os }} - name: Azure CLI Login if: github.event_name != 'pull_request' - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} @@ -546,7 +546,7 @@ jobs: os: ${{ runner.os }} - name: Azure CLI Login if: github.event_name != 'pull_request' - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} @@ -607,7 +607,7 @@ jobs: os: ${{ runner.os }} - name: Azure CLI Login if: github.event_name != 'pull_request' - uses: azure/login@7ddb5af1ef8758cf1353cf3b42f940aee27ba21c # v3.0.2 + uses: azure/login@a641126d1b8aa4d1fa005f4f92df94a3a4c4c906 # v3.1.0 with: client-id: ${{ secrets.AZURE_CLIENT_ID }} tenant-id: ${{ secrets.AZURE_TENANT_ID }} From 459fe6eca83206fde3ece88c1023fa6bc666f63e Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:06:53 +0000 Subject: [PATCH 19/26] Build(deps): Bump the codeql-actions group across 1 directory with 3 updates (#8480) Bumps the codeql-actions group with 3 updates in the / directory: [github/codeql-action/init](https://github.com/github/codeql-action), [github/codeql-action/autobuild](https://github.com/github/codeql-action) and [github/codeql-action/analyze](https://github.com/github/codeql-action). Updates `github/codeql-action/init` from 4.37.9 to 4.38.0 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/github/codeql-action/compare/cdf488f595d80d6e07e03d4674febd5ab45fa938...b96794f015dfd88f77b49b1c93e0fa7110f94c63) Updates `github/codeql-action/autobuild` from 4.37.9 to 4.38.0 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/github/codeql-action/compare/cdf488f595d80d6e07e03d4674febd5ab45fa938...b96794f015dfd88f77b49b1c93e0fa7110f94c63) Updates `github/codeql-action/analyze` from 4.37.9 to 4.38.0 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/github/codeql-action/compare/cdf488f595d80d6e07e03d4674febd5ab45fa938...b96794f015dfd88f77b49b1c93e0fa7110f94c63) --- updated-dependencies: - dependency-name: github/codeql-action/init dependency-version: 4.38.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: codeql-actions - dependency-name: github/codeql-action/autobuild dependency-version: 4.38.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: codeql-actions - dependency-name: github/codeql-action/analyze dependency-version: 4.38.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: codeql-actions ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/codeql-analysis.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/codeql-analysis.yml b/.github/workflows/codeql-analysis.yml index 62245624129..00b38da40c4 100644 --- a/.github/workflows/codeql-analysis.yml +++ b/.github/workflows/codeql-analysis.yml @@ -38,7 +38,7 @@ jobs: # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL - uses: github/codeql-action/init@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + uses: github/codeql-action/init@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 with: languages: ${{ matrix.language }} # If you wish to specify custom queries, you can do so here or in a config file. @@ -51,7 +51,7 @@ jobs: # Autobuild attempts to build any compiled languages (C/C++, C#, Go, or Java). # If this step fails, then you should remove it and run the build manually (see below) - name: Autobuild - uses: github/codeql-action/autobuild@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + uses: github/codeql-action/autobuild@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 # ℹī¸ Command-line programs to run using the OS shell. # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun @@ -64,6 +64,6 @@ jobs: # ./location_of_script_within_repo/buildscript.sh - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4.37.9 + uses: github/codeql-action/analyze@b96794f015dfd88f77b49b1c93e0fa7110f94c63 # v4.38.0 with: category: "/language:${{matrix.language}}" From 7a5acf14c363eb0fa50f697befcad7e68c2baa69 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:32:46 +0000 Subject: [PATCH 20/26] Build(deps-dev): Bump agent-framework-openai in /python/packages/lab (#8473) Bumps [agent-framework-openai](https://github.com/microsoft/agent-framework) from 1.14.2 to 1.14.3. - [Release notes](https://github.com/microsoft/agent-framework/releases) - [Commits](https://github.com/microsoft/agent-framework/commits) --- updated-dependencies: - dependency-name: agent-framework-openai dependency-version: 1.14.3 dependency-type: direct:development update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- python/packages/lab/uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/lab/uv.lock b/python/packages/lab/uv.lock index a642845cb82..daa6c9ffcb0 100644 --- a/python/packages/lab/uv.lock +++ b/python/packages/lab/uv.lock @@ -155,15 +155,15 @@ tau2 = [{ name = "tau2", git = "https://github.com/sierra-research/tau2-bench?re [[package]] name = "agent-framework-openai" -version = "1.14.2" +version = "1.14.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "agent-framework-core" }, { name = "openai" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/89/b2/c784f2c11b03f586a61775b1c5f61710379a1069257698dd803e7fbb03fa/agent_framework_openai-1.14.2.tar.gz", hash = "sha256:a09684d35bc0202b8fffea6fbd187480bc56a95027ed352d699a839d536a12ff", size = 63266, upload-time = "2026-09-03T09:53:52.239Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/7c/6ef5b4b2e5fd8ab4750c41edaca6755694a7574e4b884c3cbdffa0a1846a/agent_framework_openai-1.14.3.tar.gz", hash = "sha256:187587bb1026f0989bbfcbb1f94e058103cbfc5066006d4b82d4c00f94469a53", size = 63541, upload-time = "2026-09-10T09:29:28.72Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/43/de/6bdf51e5671b9f72ba83cc8b9ae3bc616c23bdb7f30365756c7d8baa0e74/agent_framework_openai-1.14.2-py3-none-any.whl", hash = "sha256:3da7d1f81bcd2bdbae8ba833647a4a8c2da1f1250bd3cc7947b2eb4e67e10c61", size = 68397, upload-time = "2026-09-03T09:53:41.964Z" }, + { url = "https://files.pythonhosted.org/packages/a2/d0/79a299476fb1584c04c135b0ce9f4a643d3633fd36a40effc8fbdad42da6/agent_framework_openai-1.14.3-py3-none-any.whl", hash = "sha256:b24b19b641531ef09e5cf52d5e5d20ba1be7299477d721e3516fc2da55e7f1ef", size = 68840, upload-time = "2026-09-10T09:29:06.319Z" }, ] [[package]] From 12603b63bac3ddc0cc1432ac51b211c280147c61 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:33:07 +0000 Subject: [PATCH 21/26] Build(deps): Bump huggingface-hub in /python/packages/lab (#8472) Bumps [huggingface-hub](https://github.com/huggingface/huggingface_hub) from 1.30.0 to 1.31.0. - [Release notes](https://github.com/huggingface/huggingface_hub/releases) - [Commits](https://github.com/huggingface/huggingface_hub/compare/v1.30.0...v1.31.0) --- updated-dependencies: - dependency-name: huggingface-hub dependency-version: 1.31.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- python/packages/lab/uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/lab/uv.lock b/python/packages/lab/uv.lock index daa6c9ffcb0..c18a64d328b 100644 --- a/python/packages/lab/uv.lock +++ b/python/packages/lab/uv.lock @@ -2287,7 +2287,7 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.30.0" +version = "1.31.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -2300,9 +2300,9 @@ dependencies = [ { name = "tqdm" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/35/97/2eb4abaa5b969ed385066a0496a3823b3ff467fc1082e2202955f1867d60/huggingface_hub-1.30.0.tar.gz", hash = "sha256:e6a6120bc8c8e2723d03648434ee247088cceb55ba7067e7d34d692cad5fdb57", size = 964291, upload-time = "2026-09-03T10:05:14.053Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9b/f0/61159db90b5cd275d55516fe27920828e7d3be4053fdbdb27c3f70e5f1ef/huggingface_hub-1.31.0.tar.gz", hash = "sha256:f8e9e710a210613fa5d0f26bba6da05ef4aef9fba5a0f23f508f5ac4d08b6f90", size = 968039, upload-time = "2026-09-10T10:27:22.724Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/0e/3e45bbe0dd48f4e56b1d46649d342de853cd1c7e815323472ab62687f153/huggingface_hub-1.30.0-py3-none-any.whl", hash = "sha256:96ae0a8e99a234374a6fe43e989ebd21c04640b91ab2927e7e5773ba1131ca59", size = 796795, upload-time = "2026-09-03T10:05:12.21Z" }, + { url = "https://files.pythonhosted.org/packages/c3/7f/3f886a625043b77312b80da2f2bf00b5ecbf5a73061af1aa0259cd258c9d/huggingface_hub-1.31.0-py3-none-any.whl", hash = "sha256:9dbb6a503cbe2494ea666695207e7262d410659e09134059deb83e5480864667", size = 798313, upload-time = "2026-09-10T10:27:20.798Z" }, ] [[package]] From 06f790905fa3460e13211b0a31d22a6800dbc607 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:33:23 +0000 Subject: [PATCH 22/26] Build(deps): Bump agent-framework-core in /python/packages/lab (#8471) Bumps [agent-framework-core](https://github.com/microsoft/agent-framework) from 1.17.0 to 1.18.0. - [Release notes](https://github.com/microsoft/agent-framework/releases) - [Commits](https://github.com/microsoft/agent-framework/compare/dotnet-1.17.0...dotnet-1.18.0) --- updated-dependencies: - dependency-name: agent-framework-core dependency-version: 1.18.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- python/packages/lab/uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/lab/uv.lock b/python/packages/lab/uv.lock index c18a64d328b..225d4b6a4b6 100644 --- a/python/packages/lab/uv.lock +++ b/python/packages/lab/uv.lock @@ -49,7 +49,7 @@ wheels = [ [[package]] name = "agent-framework-core" -version = "1.17.0" +version = "1.18.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "msgspec" }, @@ -58,9 +58,9 @@ dependencies = [ { name = "python-dotenv" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1d/a6/a47518b61275a7b001d3acebae6857ad98b7d10303ef8d7c9b6546b1a450/agent_framework_core-1.17.0.tar.gz", hash = "sha256:c1b69ab7f496d1ec9ab5dcb0afc5a99e60e5c820f03f5a755a9104fd89d01394", size = 593864, upload-time = "2026-09-03T09:53:46.019Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2e/6e/e0d661f1770bf3a6d70248c5964e684d2bb648b253d53f3dca315f393535/agent_framework_core-1.18.0.tar.gz", hash = "sha256:ef84e1ec570036b81a93a9e924bc4281380f6240ab9bddb8a2db7669282032b0", size = 652594, upload-time = "2026-09-10T09:29:19.938Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9c/cb/1fceefa6dc720774fc8180a8651045e583de550b946cc59c291684fcbbf6/agent_framework_core-1.17.0-py3-none-any.whl", hash = "sha256:75958ff692a38bf0c627aaa910bae6c4a89569dfec68e1e79eac8e206d88874d", size = 651715, upload-time = "2026-09-03T09:53:33.222Z" }, + { url = "https://files.pythonhosted.org/packages/98/a6/10b844bf70bd174d32635e84a5cffa2ada1df1555266e00fc5c3a402d744/agent_framework_core-1.18.0-py3-none-any.whl", hash = "sha256:75f2fac5eed229c62f0665630cf1478eb45204bde41200a4c94dab57d441cc84", size = 712930, upload-time = "2026-09-10T09:28:54.632Z" }, ] [[package]] From 5494410467000939eca8236efb7b48b51f51eb1e Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:33:42 +0000 Subject: [PATCH 23/26] Build(deps-dev): Bump pyright (#8470) Bumps the python-type-checkers group with 1 update in the /python/packages/lab directory: [pyright](https://github.com/RobertCraigie/pyright-python). Updates `pyright` from 1.1.411 to 1.1.414 - [Release notes](https://github.com/RobertCraigie/pyright-python/releases) - [Commits](https://github.com/RobertCraigie/pyright-python/compare/v1.1.411...v1.1.414) --- updated-dependencies: - dependency-name: pyright dependency-version: 1.1.414 dependency-type: direct:development update-type: version-update:semver-patch dependency-group: python-type-checkers ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- python/packages/lab/pyproject.toml | 2 +- python/packages/lab/uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/python/packages/lab/pyproject.toml b/python/packages/lab/pyproject.toml index e865740007c..c197c62e4c0 100644 --- a/python/packages/lab/pyproject.toml +++ b/python/packages/lab/pyproject.toml @@ -64,7 +64,7 @@ dev = [ "pytest-asyncio==1.4.0", "pytest-cov==7.1.0", "mypy==2.3.1", - "pyright==1.1.411", + "pyright==1.1.414", #tasks "poethepoet==0.48.0", "rich>=13.7.1,<15.0.0", diff --git a/python/packages/lab/uv.lock b/python/packages/lab/uv.lock index 225d4b6a4b6..6c2c3b52248 100644 --- a/python/packages/lab/uv.lock +++ b/python/packages/lab/uv.lock @@ -141,7 +141,7 @@ dev = [ { name = "mypy", specifier = "==2.3.1" }, { name = "poethepoet", specifier = "==0.48.0" }, { name = "prek", specifier = "==0.5.2" }, - { name = "pyright", specifier = "==1.1.411" }, + { name = "pyright", specifier = "==1.1.414" }, { name = "pytest", specifier = "==9.1.1" }, { name = "pytest-asyncio", specifier = "==1.4.0" }, { name = "pytest-cov", specifier = "==7.1.0" }, @@ -4867,15 +4867,15 @@ wheels = [ [[package]] name = "pyright" -version = "1.1.411" +version = "1.1.414" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nodeenv" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/7e/ab/265f7dc69d28113ebba19092e57b075f41543b2ed048429c5f56e2b88eac/pyright-1.1.411.tar.gz", hash = "sha256:d885a0551f2e763b089a02702174e7f4ba77548cddabc972ab86d1f7f1b0f998", size = 4112861, upload-time = "2026-06-25T02:14:06.37Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/1b/244c7b710031ada80f27e579ec20d28a2285dfc318fed0339866b1047f12/pyright-1.1.414.tar.gz", hash = "sha256:523c0a97c60da6333234955c277730c9cf4f5bd6d5399e7b7d2b0fc5d3599524", size = 4154638, upload-time = "2026-09-10T12:26:53.181Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0a/49/385be530a6a5b78d1cbcd5c2e38debc8959a2fc6bdb716f4e581002979fc/pyright-1.1.411-py3-none-any.whl", hash = "sha256:dc7c72a8e2700c55baa127554040e067041ea53ccfd50bf96308cc4291c7d5d9", size = 6181526, upload-time = "2026-06-25T02:14:04.691Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ba/18b6e682ead424ad24bcc134339ae5d1b931cd9ae260540592a058a91279/pyright-1.1.414-py3-none-any.whl", hash = "sha256:2a6b4b3298c9eec174c5ed83bd338de6eee82df2992f3e1930e6199d381be36f", size = 6225049, upload-time = "2026-09-10T12:26:51.427Z" }, ] [[package]] From fa359abd0c85f8baa406c6ec86cd9ead92631b42 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:33:57 +0000 Subject: [PATCH 24/26] Build(deps-dev): Bump uv (#8469) Bumps the basics group with 1 update in the /python/packages/lab directory: [uv](https://github.com/astral-sh/uv). Updates `uv` from 0.12.9 to 0.12.12 - [Release notes](https://github.com/astral-sh/uv/releases) - [Changelog](https://github.com/astral-sh/uv/blob/main/CHANGELOG.md) - [Commits](https://github.com/astral-sh/uv/compare/0.12.9...0.12.12) --- updated-dependencies: - dependency-name: uv dependency-version: 0.12.12 dependency-type: direct:development update-type: version-update:semver-patch dependency-group: basics ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- python/packages/lab/pyproject.toml | 2 +- python/packages/lab/uv.lock | 46 +++++++++++++++--------------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/python/packages/lab/pyproject.toml b/python/packages/lab/pyproject.toml index c197c62e4c0..000f2617e73 100644 --- a/python/packages/lab/pyproject.toml +++ b/python/packages/lab/pyproject.toml @@ -58,7 +58,7 @@ math = [ [dependency-groups] dev = [ "agent-framework-openai>=1.14.2,<2", - "uv==0.12.9", + "uv==0.12.12", "ruff==0.16.6", "pytest==9.1.1", "pytest-asyncio==1.4.0", diff --git a/python/packages/lab/uv.lock b/python/packages/lab/uv.lock index 6c2c3b52248..c43034ae57d 100644 --- a/python/packages/lab/uv.lock +++ b/python/packages/lab/uv.lock @@ -149,7 +149,7 @@ dev = [ { name = "ruff", specifier = "==0.16.6" }, { name = "tomli", specifier = "==2.4.1" }, { name = "tomli-w", specifier = "==1.2.0" }, - { name = "uv", specifier = "==0.12.9" }, + { name = "uv", specifier = "==0.12.12" }, ] tau2 = [{ name = "tau2", git = "https://github.com/sierra-research/tau2-bench?rev=5ba9e3e56db57c5e4114bf7f901291f09b2c5619" }] @@ -6414,28 +6414,28 @@ wheels = [ [[package]] name = "uv" -version = "0.12.9" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d7/42/6d2be569fd280ebd982a54ee2e3d3754fedde5cf557fd7f858eb40d161cf/uv-0.12.9.tar.gz", hash = "sha256:55b8920edb2a29eeef31e246fb60d2145985e17f51ad14b26643575b82759302", size = 7134801, upload-time = "2026-09-01T21:57:47.305Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/73/4a/208c1ed443a5024bdb83cedf1e09b322e68ccc9f7e11d1f53bb1f0690174/uv-0.12.9-py3-none-linux_armv6l.whl", hash = "sha256:97062b2160a97fcfa45f10bbbe02a9c6ba4f5a50cf4e02c0f2babb80b47c1ea8", size = 22183326, upload-time = "2026-09-01T21:57:01.443Z" }, - { url = "https://files.pythonhosted.org/packages/7a/4f/7cc3d45951bb01089bf948d91a89c5b433789df8ceeba4f4cbcce0b1ef5c/uv-0.12.9-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:0cb7e8a732d81f25d364de8d6ac69f4d00d56b581918ac242eeaeffb53977588", size = 20481767, upload-time = "2026-09-01T21:57:04.573Z" }, - { url = "https://files.pythonhosted.org/packages/5f/2b/28f22620d336069038e850ff662be979866c4944126d17a19bed738e3dc1/uv-0.12.9-py3-none-macosx_11_0_arm64.whl", hash = "sha256:0af988f800778913d36f9bf6905eb21c917dccc403bb364ad739bfd383acf491", size = 17338955, upload-time = "2026-09-01T21:57:06.884Z" }, - { url = "https://files.pythonhosted.org/packages/cc/29/649ff920dcd2373a819a4e0d9c1989ae770698e634c277e5a6fbbe19e9b4/uv-0.12.9-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:cce6e716df290b7f4a31f76ef39866e1dc006e947520a660459dea00b5d55234", size = 21598566, upload-time = "2026-09-01T21:57:09.223Z" }, - { url = "https://files.pythonhosted.org/packages/35/0d/0db7cfdd5bf11ce91c53558449f0e753d3cc578d9e94bd350c10b43f17ef/uv-0.12.9-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:00ff9f34c04055ebc3e1599ee459f5ae0d42c17cd6a1d8d029f0389e64a1419c", size = 21690729, upload-time = "2026-09-01T21:57:11.866Z" }, - { url = "https://files.pythonhosted.org/packages/69/bf/a63391a9d0f930f5ec23ff6a1daa9f7806a430fe41fb1e3bfb6cb10ef737/uv-0.12.9-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:2d4a1d4c067f9e2edf703e181a71db8fc9d7af5fcdbd0053103d0a1260e04ce6", size = 21716103, upload-time = "2026-09-01T21:57:14.407Z" }, - { url = "https://files.pythonhosted.org/packages/f5/4c/4b1fea7c22abac3297fa7fc222f3703c33ffe45aa7553fa9d8bd2b0b70bb/uv-0.12.9-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:cafa85338ce1d27b7caa4aadc83ab21f350e2cd2185462fbb10a232b6e93152e", size = 22340621, upload-time = "2026-09-01T21:57:17.109Z" }, - { url = "https://files.pythonhosted.org/packages/25/57/7ac6fd2b97ea9f608b65bb9b47d686e5d64f8142b190d3c4590d74cf8524/uv-0.12.9-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:bbdcb9db8b44d80e3fcebc4205604d52ea145404efb179958560bb2a0157743a", size = 23586367, upload-time = "2026-09-01T21:57:19.494Z" }, - { url = "https://files.pythonhosted.org/packages/ff/87/3ef6bf32db7af74e7c230ef6c06c387d9be2363dfa63eb524065031c010a/uv-0.12.9-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cfb512d0e284e9fc42c3bbe9b701f430083fca23dbb67a20daeea897fb793069", size = 23279627, upload-time = "2026-09-01T21:57:22.146Z" }, - { url = "https://files.pythonhosted.org/packages/37/4b/cd04809c7ad5149faac55160925bd67aab43f803f6c01be1c32aec7d24d9/uv-0.12.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5badfd805fd88bf99b4b4f044f6e8f762f1892cab27477f4427bb473e93dd049", size = 20056367, upload-time = "2026-09-01T21:57:24.49Z" }, - { url = "https://files.pythonhosted.org/packages/29/6b/59cddb48c47124428d329ac3ed4ce91ccd8ab757ebeed3d482115d8b82ef/uv-0.12.9-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:c0f7e2b46e4c503f9497c0f41bba3570aa233685db814cd011225118f2d72b6a", size = 19383069, upload-time = "2026-09-01T21:57:27.019Z" }, - { url = "https://files.pythonhosted.org/packages/99/a6/22e4650cc9a2d018165fda5a53ca7addedc6d4f083ab7c833acc12fbe1e1/uv-0.12.9-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:445c572202243229043abaeab2fecb1b3691477c71c79feb07f638ce34febd39", size = 22392774, upload-time = "2026-09-01T21:57:29.323Z" }, - { url = "https://files.pythonhosted.org/packages/d4/37/68166b0acb7f8ff30316bf74c0e9840c7c6cbc86ee3e236382bb107d42db/uv-0.12.9-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:ce4ad780d3cefe5a8f965f3ddf942c01e71aa37e36f4fc431e72e4ee8ce5f378", size = 22517532, upload-time = "2026-09-01T21:57:31.758Z" }, - { url = "https://files.pythonhosted.org/packages/28/ac/8801d60db574aa15fba53944e0ee2e298605bf9076a5389f5e56ed51206a/uv-0.12.9-py3-none-musllinux_1_1_i686.whl", hash = "sha256:f639f8415f08918fb4ba2022bfd94d039e01c6c1eca28583125d439f7bd635fb", size = 21498480, upload-time = "2026-09-01T21:57:34.402Z" }, - { url = "https://files.pythonhosted.org/packages/ec/47/663bbc05f26c6c518c086da5df1e38fc5b2650fa6768cb36c6418d3f9559/uv-0.12.9-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:5f9325495fa0bf3e71e6b09f72e68b0d4a000067888b48a659536a8c7191a112", size = 22889848, upload-time = "2026-09-01T21:57:37.284Z" }, - { url = "https://files.pythonhosted.org/packages/58/ce/a4b1ae2a80264231a4dd0a070ffd0fe7f2bdf1d78e99d3c0c68ce315262b/uv-0.12.9-py3-none-win32.whl", hash = "sha256:f4ab95475562884fc8e72fbbf0c3da0e025d3bb25e7348544cd02b19bf77f3e1", size = 19858143, upload-time = "2026-09-01T21:57:39.571Z" }, - { url = "https://files.pythonhosted.org/packages/f5/b2/55c7fdbdacb0b110f92804cceea449dfcd5be4354ae1e997aca700145a65/uv-0.12.9-py3-none-win_amd64.whl", hash = "sha256:871e884c3329f69c12bb7426edb9e727745c3bb2db6d9e0538ce0e1bdfb64916", size = 17991226, upload-time = "2026-09-01T21:57:42.282Z" }, - { url = "https://files.pythonhosted.org/packages/85/67/cac985492192ff4a95a76b3854143062f1b3290aca3a404a812b7aadc9cf/uv-0.12.9-py3-none-win_arm64.whl", hash = "sha256:31fddf25d756a51fe33f7831ff697a122628dd3121c37c435c39c678428f937b", size = 19472343, upload-time = "2026-09-01T21:57:44.877Z" }, +version = "0.12.12" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/77/1e/7827991def6006d84adc5113bc000956e83d190da00f836167775f276af5/uv-0.12.12.tar.gz", hash = "sha256:981a00db03357d9e9eec2b0fd128ac8b6786bf9265a56b77b9d9fc22f30fec3b", size = 7162711, upload-time = "2026-09-09T16:44:12.259Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d3/5e/705f15eaf979e50e1f6ff7238fa491196944a6137f724aaf69ab93f13fd5/uv-0.12.12-py3-none-linux_armv6l.whl", hash = "sha256:93deea45061d858760f8756fcca54a3e694978fe7f96e4987ad95ab47f72d554", size = 22185914, upload-time = "2026-09-09T16:43:29.816Z" }, + { url = "https://files.pythonhosted.org/packages/c5/49/1e5c867d088fd6641f15db3f50e07770a466f54309b818242bfd21e11469/uv-0.12.12-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:d3fa745404dd52a1763bc159073deacc5c0d501cb5e5c0baeeb53bf26f6d3482", size = 20476071, upload-time = "2026-09-09T16:43:32.473Z" }, + { url = "https://files.pythonhosted.org/packages/a9/8d/ddec2493ee930ebf32676f7b3715d75915eda45c981f6621ce3e99b69254/uv-0.12.12-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e29bf70f889e801f7dc440b51b859f005311416aa3b5c42ba9dc387c7e877c00", size = 16743115, upload-time = "2026-09-09T16:43:34.574Z" }, + { url = "https://files.pythonhosted.org/packages/77/27/b7b7486ebc998ab0893169a02d87d540c47bc6dd821de7e11156e343697a/uv-0.12.12-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:bff88c94d0b09149c1017b87b9eeca16764df37e1a6717f2b64a4104e3fb6795", size = 21666817, upload-time = "2026-09-09T16:43:36.699Z" }, + { url = "https://files.pythonhosted.org/packages/af/c1/131f9de447f8f7299ebe3f5cb5089f7c7f2a62b257d81a5cffee5bacd237/uv-0.12.12-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:75c36925481ab9a3f0f31a737636597653ca228d14152420a6047a85c6a16fb0", size = 21668514, upload-time = "2026-09-09T16:43:38.962Z" }, + { url = "https://files.pythonhosted.org/packages/12/56/94253e4f0e3806236393415b3b941625f3466c0e4091330827e3f07c27a6/uv-0.12.12-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0635b6d6601cf93d442353fa54f8dddfc3129f3f74f09721fc402983a57dc172", size = 21708940, upload-time = "2026-09-09T16:43:41.327Z" }, + { url = "https://files.pythonhosted.org/packages/30/be/cd8d050f00b11f01b59b13b4e89801f8f4bfbda599fad5858703b8b81fd3/uv-0.12.12-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e87f83293e6b3521660aa5c9aa5ab2a34fdc87845c90a466563cfbd53acf07de", size = 22402933, upload-time = "2026-09-09T16:43:43.686Z" }, + { url = "https://files.pythonhosted.org/packages/1a/84/c131548ba3e7726ad105d2022a9e761e09225736678d771cfac502141674/uv-0.12.12-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6bd1bf48fbc973c0700bf35717204b1e2a659e54fb68a010621b4949128f4a7e", size = 23701106, upload-time = "2026-09-09T16:43:46.16Z" }, + { url = "https://files.pythonhosted.org/packages/6c/0e/f9b3feda985e48d87a6c8375a1255af04af3e759dc92329387215830cd15/uv-0.12.12-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4a8535ca0ebd23d03a5a19b7c7a1c7567aac2aeaa06ee76716e133339e091c56", size = 23357681, upload-time = "2026-09-09T16:43:48.623Z" }, + { url = "https://files.pythonhosted.org/packages/be/c2/e8ff20e5f0bf1011995688784bc71ed7edbc7ccd3836e4967895680b2968/uv-0.12.12-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa5df02fc619a3cc7a58810d6ffeb80ca1e01404b8ef7239bd1cf2103c02cacf", size = 20041994, upload-time = "2026-09-09T16:43:50.98Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b6/3c3e39422e507de31270ea3f82d66c8dcf010ed35cb8559dbdfb65078413/uv-0.12.12-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:091e942f0677f88389265c9028d64108c0641593dd8993166eae71e33fc5b5aa", size = 19458417, upload-time = "2026-09-09T16:43:53.209Z" }, + { url = "https://files.pythonhosted.org/packages/fb/3a/f39ffe3cd7f3fbc33b3f62eef26274ed3f42a01ceb18296ec6fd65a36080/uv-0.12.12-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:5b22f74a63fae875baf9af5818b08901f5a7addd5440b121245e8c4c7bb84084", size = 22474163, upload-time = "2026-09-09T16:43:55.363Z" }, + { url = "https://files.pythonhosted.org/packages/aa/9c/6a0105b5102486da1ad8d0dbfb438adc61a5e5d8397757a24fd5f6d17a43/uv-0.12.12-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:0e27713b2441d9204a35e0799c6e44ffeef9978e72ea5530cafcf91f8b24986c", size = 22599041, upload-time = "2026-09-09T16:43:57.702Z" }, + { url = "https://files.pythonhosted.org/packages/c3/14/a86c94d61245f88dfd825aef961f680d6dd7ba331b4ac5254a6de44aab36/uv-0.12.12-py3-none-musllinux_1_1_i686.whl", hash = "sha256:206446685f487f21c7eff07f74513693dd53a2374934b4324ebb1a1d665b595d", size = 21570488, upload-time = "2026-09-09T16:44:00.391Z" }, + { url = "https://files.pythonhosted.org/packages/ae/c2/f81d871bb447fcf28667fda4a15b5ad6c2719ba2cd34e51b674167dfab83/uv-0.12.12-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:98d9659e080c42c289e4552d4f01b758cea544a96e3d9e5ac8565faacf84f970", size = 22996969, upload-time = "2026-09-09T16:44:02.703Z" }, + { url = "https://files.pythonhosted.org/packages/ff/5f/e1190ec0810e3d785b14f6d9c43b744a7869d3ecec3f5cda3dd47cdcf0f7/uv-0.12.12-py3-none-win32.whl", hash = "sha256:febbb31b52cd47061e496ca12e34cf083981e167d1d0ebadf4a59a0ce369d86a", size = 19441083, upload-time = "2026-09-09T16:44:04.984Z" }, + { url = "https://files.pythonhosted.org/packages/a6/b7/aebd3504d559315be6137ebfe1f6d2258e8e2f48e528d7cc799aea3441ef/uv-0.12.12-py3-none-win_amd64.whl", hash = "sha256:4a5acd0c6087970bdf210eeb1203835117f6fb8b84828604205de510f68aa6a7", size = 17654279, upload-time = "2026-09-09T16:44:07.495Z" }, + { url = "https://files.pythonhosted.org/packages/34/f5/70f1cb8ad7dba238489a86a341000c133db54aa12cdd81e67c98ebe9a59a/uv-0.12.12-py3-none-win_arm64.whl", hash = "sha256:3577343756d8bf686c92eaca89a79fb0c1b6a85c2f0b579d535db22a1b49eb4b", size = 18863039, upload-time = "2026-09-09T16:44:09.961Z" }, ] [[package]] From a08534cc7cce4ce52a66d7d38143acd743b67d0d Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Fri, 18 Sep 2026 02:51:12 +0000 Subject: [PATCH 25/26] Python: fix Cosmos memory retrieval compatibility (#8460) --- python/packages/azure-cosmos-memory/README.md | 9 +- .../_context_provider.py | 74 +++++-- .../azure-cosmos-memory/pyproject.toml | 2 +- .../tests/test_context_provider.py | 182 +++++++++++++++++- .../tests/test_emulator.py | 4 +- python/uv.lock | 8 +- 6 files changed, 249 insertions(+), 30 deletions(-) diff --git a/python/packages/azure-cosmos-memory/README.md b/python/packages/azure-cosmos-memory/README.md index db466b5ed33..e5516ae8f0c 100644 --- a/python/packages/azure-cosmos-memory/README.md +++ b/python/packages/azure-cosmos-memory/README.md @@ -311,7 +311,7 @@ memory_provider = CosmosMemoryContextProvider( credential=DefaultAzureCredential(), # Azure credential # Memory retrieval options - top_k=5, # Number of memories to retrieve + top_k=5, # Shared fact/episode result limit min_confidence=0.7, # Minimum confidence score (0.0-1.0) memory_types=["fact", "procedural"], # Types to retrieve @@ -326,16 +326,19 @@ memory_provider = CosmosMemoryContextProvider( ### Memory Types -The provider retrieves four types of memories: +The provider retrieves three types of memories: | Type | Description | Default TTL | |------|-------------|-------------| | **fact** | Declarative knowledge ("user prefers dark mode") | None | | **procedural** | Behavioral rules ("always confirm before deleting") | None | | **episodic** | Past experiences with context and outcomes | 90 days | -| **unclassified** | Memories that couldn't be confidently classified | None | Each memory has a confidence score (0.0-1.0). Use `min_confidence` to filter low-quality extractions. +With Agent Memory Toolkit 0.3.0b2 or later, facts and selected episodes are ranked together under +the single `top_k` limit. Selected procedures are compiled separately for the current task and do +not consume that ranking budget. Older supported Toolkit versions retain their generic retrieval +behavior. ### Processing Pipeline diff --git a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py index 9fa93db0598..4c2412e97ef 100644 --- a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py +++ b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_context_provider.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import inspect import logging import sys from collections.abc import Mapping, Sequence @@ -238,6 +239,12 @@ def __init__( self.memory_client = memory_client self._cosmos_endpoint = cosmos_endpoint self._foundry_endpoint = foundry_endpoint + procedural_builder = getattr(self.memory_client, "build_procedural_context", None) + self._supports_toolkit_03_retrieval = ( + "include_episodes" in inspect.signature(self.memory_client.search_cosmos).parameters + and procedural_builder is not None + and "task" in inspect.signature(procedural_builder).parameters + ) def _resolve_user_id(self, state: dict[str, Any], session: AgentSession) -> str: """Resolve the user id for memory scoping. @@ -358,26 +365,51 @@ async def before_run( # Get user_id from state or session (warns once if no stable user_id was provided) user_id = self._resolve_user_id(state, session) - # Memory search and user-summary retrieval are independent: the user summary - # provides baseline context even when no memories match the query, so a failure - # in one must not suppress the other. They get separate error handling. - try: - results = await self.memory_client.search_cosmos( - search_terms=query_text, - user_id=user_id, - top_k=self.top_k, - memory_types=[str(t) for t in self.memory_types], - min_confidence=self.min_confidence, - ) - - if results: - # Format and inject memories - memory_content = self._format_memories(results) - context.extend_messages( - self.source_id, [Message(role="user", contents=[f"{self.context_prompt}\n{memory_content}"])] + memory_sections: list[str] = [] + + # Toolkit 0.3.0b2 compiles task-aware procedures separately. Older supported + # clients keep their existing generic procedural search behavior. + search_memory_types = [ + str(memory_type) + for memory_type in self.memory_types + if not self._supports_toolkit_03_retrieval or memory_type != "procedural" + ] + if search_memory_types: + try: + search_kwargs: dict[str, Any] = { + "search_terms": query_text, + "user_id": user_id, + "top_k": self.top_k, + "memory_types": search_memory_types, + "min_confidence": self.min_confidence, + } + if self._supports_toolkit_03_retrieval: + search_kwargs["include_episodes"] = "episodic" in self.memory_types + + results = await self.memory_client.search_cosmos(**search_kwargs) + if results: + memory_sections.append(self._format_memories(results)) + except Exception as e: + logger.warning("Failed to retrieve memories: %s", e, exc_info=True) + + if self._supports_toolkit_03_retrieval and "procedural" in self.memory_types: + try: + procedural_builder = cast("Any", self.memory_client.build_procedural_context) + procedural_context = await procedural_builder( + user_id=user_id, + task=query_text, ) - except Exception as e: - logger.warning("Failed to retrieve memories: %s", e, exc_info=True) + if procedural_context and procedural_context.strip(): + memory_sections.append(procedural_context.strip()) + except Exception as e: + logger.warning("Failed to retrieve procedural context: %s", e, exc_info=True) + + if memory_sections: + memory_content = "\n".join(memory_sections) + context.extend_messages( + self.source_id, + [Message(role="user", contents=[f"{self.context_prompt}\n{memory_content}"])], + ) # Retrieve and inject user summary as untrusted context. # This is INDEPENDENT of search results - even if no memories match the query, @@ -437,7 +469,9 @@ async def after_run( # TODO(atty57): The toolkit renamed add_cosmos -> upsert_memory (same kwargs); accept either # until the declared azure-cosmos-agent-memory floor is past the rename, then inline it. - write_turn = getattr(self.memory_client, "upsert_memory", None) or self.memory_client.add_cosmos + write_turn = getattr(self.memory_client, "upsert_memory", None) + if write_turn is None: + write_turn = getattr(self.memory_client, "add_cosmos") # ruff: ignore[get-attr-with-constant] try: # Store input messages (skip empty/whitespace-only content to avoid junk turns) diff --git a/python/packages/azure-cosmos-memory/pyproject.toml b/python/packages/azure-cosmos-memory/pyproject.toml index 07e99fe3eea..2e5f8182ae8 100644 --- a/python/packages/azure-cosmos-memory/pyproject.toml +++ b/python/packages/azure-cosmos-memory/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ ] dependencies = [ "agent-framework-core>=1.13.0,<2", - "azure-cosmos-agent-memory>=0.2.0b3", + "azure-cosmos-agent-memory>=0.2.0b3,<0.4", # azure-cosmos-agent-memory depends transitively on a prompty pre-release # (prompty>=2.0.0a9, which has no stable 2.x release yet). Declaring it here as a # direct dependency makes the pre-release "explicit" so the workspace's diff --git a/python/packages/azure-cosmos-memory/tests/test_context_provider.py b/python/packages/azure-cosmos-memory/tests/test_context_provider.py index ebee3232348..3ed16fd0a09 100644 --- a/python/packages/azure-cosmos-memory/tests/test_context_provider.py +++ b/python/packages/azure-cosmos-memory/tests/test_context_provider.py @@ -13,7 +13,7 @@ pytest.importorskip("azure.cosmos.agent_memory") from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, create_autospec, patch from agent_framework import AgentResponse, Message from agent_framework._sessions import AgentSession, SessionContext @@ -30,6 +30,28 @@ _STUB_AGENT: Any = None +class Toolkit03MemoryClient: + """Signature-bearing test double for the Toolkit 0.3.0b2 retrieval contract.""" + + async def search_cosmos( + self, + search_terms: str, + *, + user_id: str, + top_k: int, + memory_types: list[str], + min_confidence: float, + include_episodes: bool = False, + ) -> list[dict[str, Any]]: + return [] + + async def build_procedural_context(self, user_id: str, task: str | None = None) -> str: + return "" + + async def get_user_summary(self, user_id: str) -> dict[str, Any] | None: + return None + + async def test_before_run_marks_cosmos_memory_used_before_empty_return() -> None: provider = object.__new__(CosmosMemoryContextProvider) context = MagicMock(spec=SessionContext) @@ -59,6 +81,16 @@ def mock_memory_client() -> AsyncMock: return mock_client +@pytest.fixture +def toolkit_03_memory_client() -> Any: + """Create a mock client that exposes the Toolkit 0.3.0b2 retrieval signatures.""" + mock_client = create_autospec(Toolkit03MemoryClient, instance=True) + mock_client.search_cosmos.return_value = [] + mock_client.build_procedural_context.return_value = "" + mock_client.get_user_summary.return_value = None + return mock_client + + # -- Initialization tests ------------------------------------------------------ @@ -270,6 +302,7 @@ async def test_retrieves_and_injects_memories(self, mock_memory_client: AsyncMoc assert call_kwargs["top_k"] == 5 assert call_kwargs["memory_types"] == ["fact", "procedural"] assert call_kwargs["min_confidence"] == 0.7 + assert "include_episodes" not in call_kwargs # Verify memories added to context assert "cosmos_memory" in ctx.context_messages @@ -280,6 +313,153 @@ async def test_retrieves_and_injects_memories(self, mock_memory_client: AsyncMoc assert "0.95" in added[0].text # type: ignore assert "0.85" in added[0].text # type: ignore + async def test_older_toolkit_preserves_generic_mixed_search(self, mock_memory_client: AsyncMock) -> None: + """Older Toolkit clients receive all selected types through generic search.""" + provider = CosmosMemoryContextProvider( + memory_client=mock_memory_client, + memory_types=["fact", "episodic", "procedural"], + ) + session = AgentSession(session_id="test-session") + ctx = SessionContext(input_messages=[Message(role="user", contents=["test"])], session_id="s1") + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + call_kwargs = mock_memory_client.search_cosmos.call_args.kwargs + assert call_kwargs["memory_types"] == ["fact", "episodic", "procedural"] + assert "include_episodes" not in call_kwargs + mock_memory_client.build_procedural_context.assert_not_awaited() + + async def test_toolkit_03_retrieves_mixed_memory_types(self, toolkit_03_memory_client: Any) -> None: + """Toolkit 0.3 uses one fact/episode search plus task-aware procedural projection.""" + toolkit_03_memory_client.search_cosmos.return_value = [ + {"content": "User prefers Python", "memory_type": "fact", "confidence": 0.95}, + {"content": "User completed ML course", "memory_type": "episodic", "confidence": 0.85}, + ] + toolkit_03_memory_client.build_procedural_context.return_value = "Always provide runnable examples." + + provider = CosmosMemoryContextProvider( + memory_client=toolkit_03_memory_client, + top_k=7, + memory_types=["fact", "episodic", "procedural"], + ) + session = AgentSession(session_id="test-session") + ctx = SessionContext( + input_messages=[Message(role="user", contents=["Help me write Python"])], + session_id="s1", + ) + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + toolkit_03_memory_client.search_cosmos.assert_awaited_once_with( + search_terms="Help me write Python", + user_id="test-session", + top_k=7, + memory_types=["fact", "episodic"], + min_confidence=0.7, + include_episodes=True, + ) + toolkit_03_memory_client.build_procedural_context.assert_awaited_once_with( + user_id="test-session", + task="Help me write Python", + ) + added = ctx.context_messages["cosmos_memory"] + assert len(added) == 1 + assert "User prefers Python" in added[0].text # type: ignore + assert "User completed ML course" in added[0].text # type: ignore + assert "Always provide runnable examples." in added[0].text # type: ignore + + async def test_toolkit_03_fact_only_disables_episodes(self, toolkit_03_memory_client: Any) -> None: + """Fact-only retrieval explicitly disables episodes and skips procedures.""" + provider = CosmosMemoryContextProvider(memory_client=toolkit_03_memory_client, memory_types=["fact"]) + session = AgentSession(session_id="test-session") + ctx = SessionContext(input_messages=[Message(role="user", contents=["test"])], session_id="s1") + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + assert toolkit_03_memory_client.search_cosmos.call_args.kwargs["memory_types"] == ["fact"] + assert toolkit_03_memory_client.search_cosmos.call_args.kwargs["include_episodes"] is False + toolkit_03_memory_client.build_procedural_context.assert_not_awaited() + + async def test_toolkit_03_procedural_only_skips_generic_search(self, toolkit_03_memory_client: Any) -> None: + """Procedural-only retrieval compiles the task context without running generic search.""" + toolkit_03_memory_client.build_procedural_context.return_value = "Use the deployment runbook." + provider = CosmosMemoryContextProvider(memory_client=toolkit_03_memory_client, memory_types=["procedural"]) + session = AgentSession(session_id="test-session") + ctx = SessionContext(input_messages=[Message(role="user", contents=["Deploy the app"])], session_id="s1") + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + toolkit_03_memory_client.search_cosmos.assert_not_awaited() + toolkit_03_memory_client.build_procedural_context.assert_awaited_once_with( + user_id="test-session", + task="Deploy the app", + ) + assert "Use the deployment runbook." in ctx.context_messages["cosmos_memory"][0].text # type: ignore + + async def test_toolkit_03_empty_procedural_context_is_not_injected(self, toolkit_03_memory_client: Any) -> None: + """An empty compiled procedure does not add a context message.""" + provider = CosmosMemoryContextProvider(memory_client=toolkit_03_memory_client, memory_types=["procedural"]) + session = AgentSession(session_id="test-session") + ctx = SessionContext(input_messages=[Message(role="user", contents=["test"])], session_id="s1") + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + assert "cosmos_memory" not in ctx.context_messages + + async def test_toolkit_03_search_failure_does_not_block_other_context(self, toolkit_03_memory_client: Any) -> None: + """Generic search failure does not suppress procedures or the user summary.""" + toolkit_03_memory_client.search_cosmos.side_effect = Exception("search boom") + toolkit_03_memory_client.build_procedural_context.return_value = "Use the recovery procedure." + toolkit_03_memory_client.get_user_summary.return_value = {"content": "Prefers concise answers"} + provider = CosmosMemoryContextProvider( + memory_client=toolkit_03_memory_client, + memory_types=["fact", "procedural"], + ) + session = AgentSession(session_id="test-session") + ctx = SessionContext(input_messages=[Message(role="user", contents=["test"])], session_id="s1") + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + added = ctx.context_messages["cosmos_memory"] + assert any("Use the recovery procedure." in message.text for message in added) # type: ignore + assert any("Prefers concise answers" in message.text for message in added) # type: ignore + + async def test_toolkit_03_procedural_failure_does_not_block_other_context( + self, toolkit_03_memory_client: Any + ) -> None: + """Procedural projection failure does not suppress facts or the user summary.""" + toolkit_03_memory_client.search_cosmos.return_value = [ + {"content": "User likes hiking", "memory_type": "fact", "confidence": 0.9} + ] + toolkit_03_memory_client.build_procedural_context.side_effect = Exception("procedure boom") + toolkit_03_memory_client.get_user_summary.return_value = {"content": "Prefers concise answers"} + provider = CosmosMemoryContextProvider( + memory_client=toolkit_03_memory_client, + memory_types=["fact", "procedural"], + ) + session = AgentSession(session_id="test-session") + ctx = SessionContext(input_messages=[Message(role="user", contents=["test"])], session_id="s1") + + await provider.before_run( + agent=_STUB_AGENT, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}) + ) + + added = ctx.context_messages["cosmos_memory"] + assert any("User likes hiking" in message.text for message in added) # type: ignore + assert any("Prefers concise answers" in message.text for message in added) # type: ignore + async def test_user_summary_injected_as_untrusted_message(self, mock_memory_client: AsyncMock) -> None: """User summary is injected as an untrusted context message, not as agent instructions.""" mock_memory_client.search_cosmos.return_value = [] diff --git a/python/packages/azure-cosmos-memory/tests/test_emulator.py b/python/packages/azure-cosmos-memory/tests/test_emulator.py index caad815db4d..771a82ad802 100644 --- a/python/packages/azure-cosmos-memory/tests/test_emulator.py +++ b/python/packages/azure-cosmos-memory/tests/test_emulator.py @@ -183,7 +183,9 @@ async def test_before_run_retrieves_seeded_fact(self, emulator_provider: CosmosM # embeddings client). This lands in the memories container under the quantizedFlat # vector index, without needing LLM extraction. assert provider.memory_client is not None - seed = getattr(provider.memory_client, "upsert_memory", None) or provider.memory_client.add_cosmos + seed = getattr(provider.memory_client, "upsert_memory", None) + if seed is None: + seed = getattr(provider.memory_client, "add_cosmos") # noqa: B009 await seed( user_id=user_id, thread_id=thread_id, diff --git a/python/uv.lock b/python/uv.lock index 2ae5de419bd..047cbf03701 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -333,7 +333,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "agent-framework-core", editable = "packages/core" }, - { name = "azure-cosmos-agent-memory", specifier = ">=0.2.0b3" }, + { name = "azure-cosmos-agent-memory", specifier = ">=0.2.0b3,<0.4" }, { name = "prompty", specifier = ">=2.0.0a9" }, ] @@ -1469,7 +1469,7 @@ wheels = [ [[package]] name = "azure-cosmos-agent-memory" -version = "0.3.0b1" +version = "0.3.0b2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -1482,9 +1482,9 @@ dependencies = [ { name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/71/85/f66f7e5989125428eba72917224fbdb4fb22671de3a39ec8d24446f93b2b/azure_cosmos_agent_memory-0.3.0b1.tar.gz", hash = "sha256:4dd017bd4c69895f3e053511c549ee8b215e5eb98f684152cf852ebbb69199b3", size = 169559, upload-time = "2026-07-24T21:37:19.433Z" } +sdist = { url = "https://files.pythonhosted.org/packages/59/98/1bba425448cfe620e27818c5cc770a5c5ca69d76b8e5989c0ff69daef01d/azure_cosmos_agent_memory-0.3.0b2.tar.gz", hash = "sha256:5f3d4590b74fd15da2973d09217d9eefc2a3787c03d283210b875b05f06e68e3", size = 189529, upload-time = "2026-08-17T23:09:33.819Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/f5/18be23d17aee288562adb9ef66826829524077e7024e2b5535615c82f621/azure_cosmos_agent_memory-0.3.0b1-py3-none-any.whl", hash = "sha256:7d36e28cba2adfc1567a94d8b057c018e7195bf3d5cec14e92747cb23b27465b", size = 187471, upload-time = "2026-07-24T21:37:18.268Z" }, + { url = "https://files.pythonhosted.org/packages/e0/b4/d18966a12e3ba2c99dca4a3888bd0b35c1d1fd2c341690bd926e80550018/azure_cosmos_agent_memory-0.3.0b2-py3-none-any.whl", hash = "sha256:2af5d6e30b485052f52abc8980be980632e2ace5750eb6329e2af8ca66c5aaa9", size = 216786, upload-time = "2026-08-17T23:09:32.666Z" }, ] [[package]] From 2e065314a56e39dc53c1ee63003c7086d81b0fe7 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 18 Sep 2026 10:07:14 +0200 Subject: [PATCH 26/26] ci: update Python dependencies across projects --- .github/dependabot.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 2664f71d40b..824f537f139 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -61,7 +61,9 @@ updates: # TODO: Remove these Python Dependabot entries after we have confidence in the # Python dependency-maintenance workflow. - package-ecosystem: "pip" - directory: "python/" + directories: + - "python/" + - "python/packages/lab/" schedule: interval: "weekly" day: "thursday"