From a0bfb53caf2e62555d7a228f1dc8d3e33c8da528 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 23 Jul 2026 14:42:18 +0200 Subject: [PATCH 1/8] Python: add first-pass feature usage telemetry Add the 128-bit feature accumulator, package-local indexes, activation markers, and destination-scoped User-Agent emission for the initial Python implementation slice. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- python/packages/core/AGENTS.md | 1 + python/packages/core/README.md | 11 ++ .../packages/core/agent_framework/_agents.py | 2 + .../core/agent_framework/_compaction.py | 2 + .../core/agent_framework/_harness/_agent.py | 2 + .../_harness/_background_agents.py | 2 + .../agent_framework/_harness/_file_access.py | 2 + .../core/agent_framework/_harness/_memory.py | 2 + .../core/agent_framework/_harness/_mode.py | 2 + .../core/agent_framework/_harness/_todo.py | 2 + .../_harness/_tool_approval.py | 2 + python/packages/core/agent_framework/_mcp.py | 2 + .../core/agent_framework/_sessions.py | 5 + .../packages/core/agent_framework/_skills.py | 5 + .../core/agent_framework/_telemetry.py | 78 ++++++++ .../_workflows/_workflow_builder.py | 7 +- .../core/tests/core/test_telemetry.py | 176 ++++++++++++++++++ .../foundry/agent_framework_foundry/_agent.py | 17 +- .../agent_framework_foundry/_chat_client.py | 13 +- .../_embedding_client.py | 21 ++- .../agent_framework_foundry/_feature_usage.py | 63 +++++++ .../agent_framework_foundry/_foundry_evals.py | 32 +++- .../_memory_provider.py | 9 +- .../tests/foundry/test_foundry_agent.py | 12 +- .../tests/foundry/test_foundry_chat_client.py | 22 ++- .../foundry/test_foundry_embedding_client.py | 19 +- .../foundry/test_foundry_memory_provider.py | 3 +- .../_feature_usage.py | 9 + .../_foundry_local_client.py | 6 +- .../tests/test_foundry_local_client.py | 5 + .../agent_framework_openai/_chat_client.py | 21 ++- .../_chat_completion_client.py | 22 ++- .../_embedding_client.py | 6 +- .../agent_framework_openai/_feature_usage.py | 44 +++++ .../openai/agent_framework_openai/_shared.py | 4 + .../test_openai_chat_completion_client.py | 37 ++++ .../openai/test_openai_embedding_client.py | 21 +++ .../openai/tests/openai/test_openai_shared.py | 31 +++ .../_concurrent.py | 5 +- .../_feature_usage.py | 13 ++ .../_group_chat.py | 5 +- .../_handoff.py | 5 +- .../_magentic.py | 5 +- .../_orchestration_request_info.py | 3 +- .../_sequential.py | 5 +- .../_workflow_builder.py | 12 ++ .../orchestrations/tests/test_magentic.py | 15 ++ .../orchestrations/tests/test_sequential.py | 18 ++ python/samples/README.md | 2 + 49 files changed, 766 insertions(+), 42 deletions(-) create mode 100644 python/packages/foundry/agent_framework_foundry/_feature_usage.py create mode 100644 python/packages/foundry_local/agent_framework_foundry_local/_feature_usage.py create mode 100644 python/packages/openai/agent_framework_openai/_feature_usage.py create mode 100644 python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py create mode 100644 python/packages/orchestrations/agent_framework_orchestrations/_workflow_builder.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index ecbb455e140..b8134d3cd9e 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -17,6 +17,7 @@ agent_framework/ ├── _sessions.py # AgentSession and context provider abstractions ├── _skills.py # Agent Skills system (models, executors, provider) ├── _mcp.py # Model Context Protocol support +├── _telemetry.py # User-Agent identity and internal feature-usage mask ├── _workflows/ # Workflow orchestration (sequential, concurrent, handoff, etc.) ├── openai/ # Built-in OpenAI client ├── azure/ # Lazy-loading entry point for Azure integrations diff --git a/python/packages/core/README.md b/python/packages/core/README.md index 5ac622551f7..cc145b0aa3b 100644 --- a/python/packages/core/README.md +++ b/python/packages/core/README.md @@ -53,6 +53,17 @@ client = OpenAIChatClient( ) ``` +### Telemetry controls + +Agent Framework adds its package/version User-Agent to supported client +requests. Approved Microsoft Foundry and Azure OpenAI request paths can also +carry a documented feature-usage token. + +- `AGENT_FRAMEWORK_FEATURE_MASK_DISABLED=true` disables only the feature-usage + token while retaining the package/version User-Agent. +- `AGENT_FRAMEWORK_USER_AGENT_DISABLED=true` disables the entire Agent Framework + User-Agent contribution, including the feature token. + See the following [getting started samples](https://github.com/microsoft/agent-framework/tree/main/python/samples/01-get-started) for more information. ## 2. Create a Simple Agent diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c6d20f85abb..3ded2f67639 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -37,6 +37,7 @@ SessionContext, is_local_history_conversation_id, ) +from ._telemetry import FeatureIndex, mark_feature_used from ._types import ( AgentResponse, AgentResponseUpdate, @@ -1773,6 +1774,7 @@ def run( client_kwargs: Mapping[str, Any] | None = None, ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: """Run the agent.""" + mark_feature_used(FeatureIndex.CORE_AGENT) super_run = cast( "Callable[..., Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]]", super().run, diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 250ee98e9be..c56657febbc 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -17,6 +17,7 @@ ) from ._sessions import ContextProvider +from ._telemetry import FeatureIndex, mark_feature_used from ._types import ChatResponse, Content, Message if TYPE_CHECKING: @@ -1539,6 +1540,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Compact messages already present in the context from earlier providers.""" + mark_feature_used(FeatureIndex.CORE_COMPACTION_PROVIDER) if self.before_strategy is None: return diff --git a/python/packages/core/agent_framework/_harness/_agent.py b/python/packages/core/agent_framework/_harness/_agent.py index 7e17a3d3509..870b8a493bf 100644 --- a/python/packages/core/agent_framework/_harness/_agent.py +++ b/python/packages/core/agent_framework/_harness/_agent.py @@ -22,6 +22,7 @@ from .._feature_stage import ExperimentalFeature, warn_experimental_feature from .._sessions import ContextProvider, HistoryProvider, InMemoryHistoryProvider, MessageInjectionMiddleware from .._skills import SkillsProvider +from .._telemetry import FeatureIndex, mark_feature_used from .._types import ChatOptions from ._background_agents import BackgroundAgentsProvider from ._file_access import AgentFileStore, FileAccessProvider, FileSystemAgentFileStore @@ -674,5 +675,6 @@ def create_harness_agent( # Set the telemetry provider name after construction. agent.otel_provider_name = otel_provider_name or HARNESS_AGENT_PROVIDER_NAME + mark_feature_used(FeatureIndex.CORE_HARNESS_AGENT) return agent diff --git a/python/packages/core/agent_framework/_harness/_background_agents.py b/python/packages/core/agent_framework/_harness/_background_agents.py index 22fac3bbd44..c1c70b64a34 100644 --- a/python/packages/core/agent_framework/_harness/_background_agents.py +++ b/python/packages/core/agent_framework/_harness/_background_agents.py @@ -19,6 +19,7 @@ from .._feature_stage import ExperimentalFeature, experimental from .._serialization import SerializationMixin from .._sessions import AgentSession, ContextProvider, SessionContext +from .._telemetry import FeatureIndex, mark_feature_used from .._tools import tool from .._types import AgentResponse, Message @@ -320,6 +321,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject background agent tools and instructions before the model runs.""" + mark_feature_used(FeatureIndex.CORE_BACKGROUND_AGENTS_PROVIDER) del agent, state provider_state = _get_provider_state(session, source_id=self.source_id) diff --git a/python/packages/core/agent_framework/_harness/_file_access.py b/python/packages/core/agent_framework/_harness/_file_access.py index c0c3b15ea73..b70d78081ed 100644 --- a/python/packages/core/agent_framework/_harness/_file_access.py +++ b/python/packages/core/agent_framework/_harness/_file_access.py @@ -38,6 +38,7 @@ from .._feature_stage import ExperimentalFeature, experimental from .._serialization import SerializationMixin from .._sessions import AgentSession, ContextProvider, SessionContext +from .._telemetry import FeatureIndex, mark_feature_used from .._tools import ApprovalMode, tool from .._types import Content @@ -1454,6 +1455,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject file-access tools and instructions before the model runs.""" + mark_feature_used(FeatureIndex.CORE_FILE_ACCESS_PROVIDER) readonly_approval: ApprovalMode = "never_require" if self.disable_readonly_tool_approval else "always_require" write_approval: ApprovalMode = "never_require" if self.disable_write_tool_approval else "always_require" diff --git a/python/packages/core/agent_framework/_harness/_memory.py b/python/packages/core/agent_framework/_harness/_memory.py index e27abb05d2b..6b4fe878e17 100644 --- a/python/packages/core/agent_framework/_harness/_memory.py +++ b/python/packages/core/agent_framework/_harness/_memory.py @@ -20,6 +20,7 @@ from .._compaction import group_messages from .._feature_stage import ExperimentalFeature, experimental from .._sessions import AgentSession, FileHistoryProvider, HistoryProvider, JsonDumps, JsonLoads, SessionContext +from .._telemetry import FeatureIndex, mark_feature_used from .._tools import tool from .._types import ChatResponse, Message from ..exceptions import ChatClientException @@ -1170,6 +1171,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject ``MEMORY.md`` and selected topic files before the model runs.""" + mark_feature_used(FeatureIndex.CORE_MEMORY_PROVIDER) state.clear() state.update(self.store.export_provider_state(session)) diff --git a/python/packages/core/agent_framework/_harness/_mode.py b/python/packages/core/agent_framework/_harness/_mode.py index a3ab0ad6205..11cc018b178 100644 --- a/python/packages/core/agent_framework/_harness/_mode.py +++ b/python/packages/core/agent_framework/_harness/_mode.py @@ -7,6 +7,7 @@ from typing import Any, cast from .._sessions import AgentSession, ContextProvider, SessionContext +from .._telemetry import FeatureIndex, mark_feature_used from .._tools import tool from .._types import Message @@ -272,6 +273,7 @@ async def before_run( context: The session context to receive instructions and tools. state: Per-provider invocation state. """ + mark_feature_used(FeatureIndex.CORE_AGENT_MODE_PROVIDER) del agent, state current_mode = get_agent_mode( session, diff --git a/python/packages/core/agent_framework/_harness/_todo.py b/python/packages/core/agent_framework/_harness/_todo.py index 4dcec64d6f8..1c675097326 100644 --- a/python/packages/core/agent_framework/_harness/_todo.py +++ b/python/packages/core/agent_framework/_harness/_todo.py @@ -17,6 +17,7 @@ from .._feature_stage import ExperimentalFeature, experimental from .._serialization import SerializationMixin from .._sessions import AgentSession, ContextProvider, SessionContext +from .._telemetry import FeatureIndex, mark_feature_used from .._tools import tool from .._types import Message @@ -498,6 +499,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject todo tools and instructions before the model runs.""" + mark_feature_used(FeatureIndex.CORE_TODO_PROVIDER) del agent, state @tool(name="todos_add", approval_mode="never_require") diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index ce34ac271e0..390c516ce1b 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -12,6 +12,7 @@ from .._middleware import AgentContext, AgentMiddleware from .._serialization import SerializationMixin from .._sessions import AgentSession +from .._telemetry import FeatureIndex, mark_feature_used from .._types import ( AgentResponse, AgentResponseUpdate, @@ -379,6 +380,7 @@ def __init__( async def process(self, context: AgentContext, call_next: Callable[[], Awaitable[None]]) -> None: """Process one agent invocation.""" + mark_feature_used(FeatureIndex.CORE_TOOL_APPROVAL) if context.session is None: raise RuntimeError("ToolApprovalMiddleware requires an AgentSession.") diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 9071deafcf3..1ff60a53c83 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -29,6 +29,7 @@ _warn_on_feature_use, # pyright: ignore[reportPrivateUsage] experimental, ) +from ._telemetry import FeatureIndex, mark_feature_used from ._tools import FunctionTool from ._types import ( ChatOptions, @@ -1268,6 +1269,7 @@ async def _reconnect_without_loading(self) -> None: await self._run_on_lifecycle_owner("connect", reset=True, load_configured=False) async def connect(self, *, reset: bool = False) -> None: + mark_feature_used(FeatureIndex.CORE_MCP) if self._is_lifecycle_owner_task(): await self._connect_on_owner(reset=reset) return diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 7010606bef1..e4fec690457 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -30,6 +30,7 @@ from ._feature_stage import ExperimentalFeature, experimental from ._middleware import ChatContext, ChatMiddleware +from ._telemetry import FeatureIndex, mark_feature_used from ._types import ( AgentResponse, AgentRunInputs, @@ -1168,6 +1169,7 @@ async def get_messages( self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any ) -> list[Message]: """Retrieve messages from session state.""" + mark_feature_used(FeatureIndex.CORE_IN_MEMORY_HISTORY_PROVIDER) if state is None: return [] messages = list(state.get("messages", [])) @@ -1184,6 +1186,7 @@ async def save_messages( **kwargs: Any, ) -> None: """Persist messages to session state.""" + mark_feature_used(FeatureIndex.CORE_IN_MEMORY_HISTORY_PROVIDER) if state is None: return existing = state.get("messages", []) @@ -1303,6 +1306,7 @@ async def get_messages( **kwargs: Any, ) -> list[Message]: """Retrieve messages from the session's JSON Lines file.""" + mark_feature_used(FeatureIndex.CORE_FILE_HISTORY_PROVIDER) del state, kwargs file_path = self._session_file_path(session_id) async_lock = self._session_async_write_lock(file_path) @@ -1354,6 +1358,7 @@ async def save_messages( **kwargs: Any, ) -> None: """Append messages to the session's JSON Lines file.""" + mark_feature_used(FeatureIndex.CORE_FILE_HISTORY_PROVIDER) del state, kwargs if not messages: return diff --git a/python/packages/core/agent_framework/_skills.py b/python/packages/core/agent_framework/_skills.py index 388a6417369..b3fdbf4d1bf 100644 --- a/python/packages/core/agent_framework/_skills.py +++ b/python/packages/core/agent_framework/_skills.py @@ -62,6 +62,7 @@ from ._feature_stage import ExperimentalFeature, experimental from ._sessions import ContextProvider +from ._telemetry import FeatureIndex, mark_feature_used from ._tools import ApprovalMode, FunctionTool if TYPE_CHECKING: @@ -2407,6 +2408,7 @@ async def before_run( context: Session context to extend with instructions and tools. state: Mutable per-run state dictionary (unused by this provider). """ + mark_feature_used(FeatureIndex.CORE_SKILLS_PROVIDER) source_context = SkillsSourceContext(agent=agent, session=session) skills, instructions, tools = await self._create_context(source_context) @@ -2886,6 +2888,7 @@ async def get_skills(self, context: SkillsSourceContext) -> list[Skill]: Returns: A list of discovered file-based skills. """ + mark_feature_used(FeatureIndex.CORE_FILE_SKILLS_SOURCE) skills: dict[str, FileSkill] = {} discovered = FileSkillsSource._discover_skill_directories(self._skill_paths) @@ -3591,6 +3594,7 @@ async def get_skills(self, context: SkillsSourceContext) -> list[Skill]: Returns: A list of :class:`Skill` instances. """ + mark_feature_used(FeatureIndex.CORE_IN_MEMORY_SKILLS_SOURCE) return self._skills @@ -4374,6 +4378,7 @@ async def get_skills(self, context: SkillsSourceContext) -> list[Skill]: Returns: A list of discovered :class:`MCPSkill` instances. """ + mark_feature_used(FeatureIndex.CORE_MCP_SKILLS_SOURCE) index = await self._try_read_index() if index is None: return [] diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index 89b23927f9d..a7c48f3b48f 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -5,6 +5,9 @@ import contextlib import logging import os +import re +import threading +from enum import IntEnum from typing import Any, Final from . import __version__ as version_info @@ -15,6 +18,8 @@ # Note that if this environment variable does not exist, user agent telemetry is enabled. USER_AGENT_TELEMETRY_DISABLED_ENV_VAR = "AGENT_FRAMEWORK_USER_AGENT_DISABLED" IS_TELEMETRY_ENABLED = os.environ.get(USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, "false").lower() not in ["true", "1"] +FEATURE_MASK_DISABLED_ENV_VAR = "AGENT_FRAMEWORK_FEATURE_MASK_DISABLED" +FEATURE_REGISTRY_VERSION = 1 APP_INFO = ( { @@ -27,6 +32,29 @@ HTTP_USER_AGENT: Final[str] = "agent-framework-python" AGENT_FRAMEWORK_USER_AGENT = f"{HTTP_USER_AGENT}/{version_info}" + +class FeatureIndex(IntEnum): + """Core-owned indexes in the Python feature-usage registry.""" + + CORE_AGENT = 0 + CORE_HARNESS_AGENT = 1 + CORE_WORKFLOW = 2 + CORE_MCP = 3 + CORE_TOOL_APPROVAL = 4 + CORE_MEMORY_PROVIDER = 5 + CORE_SKILLS_PROVIDER = 6 + CORE_FILE_ACCESS_PROVIDER = 7 + CORE_COMPACTION_PROVIDER = 8 + CORE_TODO_PROVIDER = 9 + CORE_AGENT_MODE_PROVIDER = 10 + CORE_BACKGROUND_AGENTS_PROVIDER = 11 + CORE_IN_MEMORY_HISTORY_PROVIDER = 12 + CORE_FILE_HISTORY_PROVIDER = 13 + CORE_FILE_SKILLS_SOURCE = 14 + CORE_IN_MEMORY_SKILLS_SOURCE = 15 + CORE_MCP_SKILLS_SOURCE = 16 + + # This environment variable is reserved by the Foundry hosting environment to # indicate that the agent is running in a hosted environment. _FOUNDRY_HOSTING_ENV_VAR = "FOUNDRY_HOSTING_ENVIRONMENT" @@ -35,6 +63,9 @@ _user_agent_prefixes: set[str] = set() _hosted_env_detected: bool = False +_feature_mask = 0 +_feature_mask_lock = threading.Lock() +_feature_comment_pattern = re.compile(r"\s+\(feat=v\d+\.[0-9a-fA-F]+\)") def _add_user_agent_prefix(prefix: str) -> None: @@ -96,6 +127,53 @@ def get_user_agent() -> str: return f"{'/'.join(sorted(_user_agent_prefixes))}/{AGENT_FRAMEWORK_USER_AGENT}" +def _feature_mask_enabled() -> bool: + """Return whether feature-usage marking and emission are enabled.""" + return IS_TELEMETRY_ENABLED and os.environ.get(FEATURE_MASK_DISABLED_ENV_VAR, "false").lower() not in ("true", "1") + + +def mark_feature_used(index: IntEnum | int) -> None: + """Mark a feature as used in the process-global feature mask.""" + if not _feature_mask_enabled(): + return + + feature_index = int(index) + if not 0 <= feature_index < 128: + raise ValueError(f"Feature index must be in range 0..127, got {feature_index}") + + global _feature_mask + with _feature_mask_lock: + _feature_mask |= 1 << feature_index + + +def get_feature_token() -> str | None: + """Return the current versioned feature token, or None when empty or disabled.""" + if not _feature_mask_enabled(): + return None + + with _feature_mask_lock: + feature_mask = _feature_mask + if feature_mask == 0: + return None + return f"v{FEATURE_REGISTRY_VERSION}.{feature_mask:x}" + + +def apply_feature_token(user_agent: str) -> str: + """Append or refresh the live feature token in a User-Agent value.""" + base_user_agent = remove_feature_token(user_agent) + token = get_feature_token() + if token is None: + return base_user_agent + if not base_user_agent: + return f"(feat={token})" + return f"{base_user_agent} (feat={token})" + + +def remove_feature_token(user_agent: str) -> str: + """Remove the Agent Framework feature token from a User-Agent value.""" + return _feature_comment_pattern.sub("", user_agent).strip() + + def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None) -> dict[str, Any]: """Prepend "agent-framework" to the User-Agent in the headers. diff --git a/python/packages/core/agent_framework/_workflows/_workflow_builder.py b/python/packages/core/agent_framework/_workflows/_workflow_builder.py index de389b3eb58..5d2584c1d8a 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_builder.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_builder.py @@ -4,9 +4,10 @@ import sys import uuid from collections.abc import Callable, Sequence -from typing import Any, Literal +from typing import Any, ClassVar, Literal from .._agents import SupportsAgentRun +from .._telemetry import FeatureIndex, mark_feature_used from ..observability import OtelAttr, capture_exception, create_workflow_span from ._agent_executor import AgentExecutor from ._agent_utils import resolve_agent_id @@ -85,6 +86,8 @@ async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None: print(events.get_outputs()) # ['OLLEH'] """ + _FEATURE_USAGE_INDEX: ClassVar[FeatureIndex | None] = FeatureIndex.CORE_WORKFLOW + def __init__( self, max_iterations: int = DEFAULT_MAX_ITERATIONS, @@ -800,6 +803,8 @@ async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None: events = await workflow.run("hello") print(events.get_outputs()) # outputs from planner and answerer """ + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) # Create workflow build span that includes validation and workflow creation with create_workflow_span(OtelAttr.WORKFLOW_BUILD_SPAN) as span: try: diff --git a/python/packages/core/tests/core/test_telemetry.py b/python/packages/core/tests/core/test_telemetry.py index b0b01706ef1..c405c930302 100644 --- a/python/packages/core/tests/core/test_telemetry.py +++ b/python/packages/core/tests/core/test_telemetry.py @@ -1,8 +1,14 @@ # Copyright (c) Microsoft. All rights reserved. +import ast +import concurrent.futures import os +import re +from pathlib import Path from unittest.mock import MagicMock, patch +import pytest + import agent_framework._telemetry as _telemetry_mod from agent_framework import ( AGENT_FRAMEWORK_USER_AGENT, @@ -13,8 +19,14 @@ from agent_framework._telemetry import ( _FOUNDRY_HOSTING_ENV_VAR, _HOSTED_USER_AGENT_PREFIX, + FEATURE_MASK_DISABLED_ENV_VAR, + FEATURE_REGISTRY_VERSION, + FeatureIndex, _add_user_agent_prefix, _detect_hosted_environment, + apply_feature_token, + get_feature_token, + mark_feature_used, ) # region Test constants @@ -35,6 +47,170 @@ def test_agent_framework_user_agent_format(): assert AGENT_FRAMEWORK_USER_AGENT.startswith("agent-framework-python/") +def _reset_feature_mask() -> None: + with _telemetry_mod._feature_mask_lock: + _telemetry_mod._feature_mask = 0 + + +def test_feature_mask_disabled_env_var() -> None: + assert FEATURE_MASK_DISABLED_ENV_VAR == "AGENT_FRAMEWORK_FEATURE_MASK_DISABLED" + + +def test_feature_registry_version() -> None: + assert FEATURE_REGISTRY_VERSION == 1 + + +def test_mark_feature_used_accumulates_and_deduplicates() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + ): + mark_feature_used(FeatureIndex.CORE_AGENT) + mark_feature_used(FeatureIndex.CORE_AGENT) + mark_feature_used(FeatureIndex.CORE_WORKFLOW) + + assert get_feature_token() == "v1.5" + + +def test_mark_feature_used_supports_bit_127() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + ): + mark_feature_used(127) + + assert get_feature_token() == f"v1.{1 << 127:x}" + + +@pytest.mark.parametrize("bit", [-1, 128]) +def test_mark_feature_used_rejects_out_of_range_bit(bit: int) -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + pytest.raises(ValueError, match="Feature index must be in range 0..127"), + ): + mark_feature_used(bit) + + +@pytest.mark.parametrize("disabled_value", ["true", "TRUE", "1"]) +def test_feature_mask_env_var_disables_marking(disabled_value: str) -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: disabled_value}), + ): + mark_feature_used(FeatureIndex.CORE_AGENT) + + assert get_feature_token() is None + + +def test_user_agent_env_var_disables_feature_mask() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", False), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + ): + mark_feature_used(FeatureIndex.CORE_AGENT) + + assert get_feature_token() is None + + +def test_apply_feature_token_adds_and_refreshes_live_mask() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + ): + mark_feature_used(FeatureIndex.CORE_AGENT) + user_agent = apply_feature_token("foundry-hosting/agent-framework-python/1.0") + assert user_agent == "foundry-hosting/agent-framework-python/1.0 (feat=v1.1)" + + mark_feature_used(FeatureIndex.CORE_WORKFLOW) + assert apply_feature_token(user_agent) == "foundry-hosting/agent-framework-python/1.0 (feat=v1.5)" + + +def test_apply_feature_token_preserves_unrelated_comments() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + ): + mark_feature_used(FeatureIndex.CORE_AGENT) + + assert apply_feature_token("agent-framework-python/1.0 (custom=value)") == ( + "agent-framework-python/1.0 (custom=value) (feat=v1.1)" + ) + + +def test_apply_feature_token_removes_stale_token_when_disabled() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "true"}), + ): + assert apply_feature_token("agent-framework-python/1.0 (feat=v1.5)") == "agent-framework-python/1.0" + + +def test_mark_feature_used_is_thread_safe() -> None: + _reset_feature_mask() + with ( + patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True), + patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "false"}), + concurrent.futures.ThreadPoolExecutor() as executor, + ): + list(executor.map(mark_feature_used, range(128))) + + assert get_feature_token() == f"v1.{(1 << 128) - 1:x}" + + +def test_declared_feature_indexes_do_not_overlap() -> None: + registry_path = next( + ( + parent / "docs" / "specs" / "feature-usage-bit-registry.md" + for parent in Path(__file__).resolve().parents + if (parent / "docs" / "specs" / "feature-usage-bit-registry.md").exists() + ), + None, + ) + if registry_path is None: + pytest.skip("Feature-usage registry is not available outside a repository checkout.") + + repository_root = registry_path.parents[2] + declaration_files = [ + repository_root / "python" / "packages" / "core" / "agent_framework" / "_telemetry.py", + *repository_root.glob("python/packages/*/agent_framework*/_feature_usage.py"), + ] + declarations: dict[int, str] = {} + for declaration_file in declaration_files: + tree = ast.parse(declaration_file.read_text(encoding="utf-8")) + for node in tree.body: + if not isinstance(node, ast.ClassDef) or node.name != "FeatureIndex": + continue + for member in node.body: + if not isinstance(member, ast.Assign) or len(member.targets) != 1: + continue + target = member.targets[0] + if not isinstance(target, ast.Name) or not isinstance(member.value, ast.Constant): + continue + index = member.value.value + if not isinstance(index, int): + continue + assert 0 <= index < 128 + assert index not in declarations, ( + f"Feature index {index} overlaps between {declarations[index]} and " + f"{declaration_file.relative_to(repository_root)}:{target.id}" + ) + declarations[index] = f"{declaration_file.relative_to(repository_root)}:{target.id}" + + registry_text = registry_path.read_text(encoding="utf-8") + python_table = registry_text.split("## Index table — Python", 1)[1].split("## Index table — .NET", 1)[0] + registry_indexes = {int(index) for index in re.findall(r"^\| (\d+) \| `[^`]+` \|", python_table, re.MULTILINE)} + assert declarations.keys() <= registry_indexes + + def test_app_info_when_telemetry_enabled(): """Test that APP_INFO is set when telemetry is enabled.""" with patch("agent_framework._telemetry.IS_TELEMETRY_ENABLED", True): diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index f93ee54aae3..21a7e28fa79 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -29,7 +29,7 @@ load_settings, ) from agent_framework._compaction import CompactionStrategy, TokenizerProtocol -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import IS_TELEMETRY_ENABLED, get_user_agent from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient from azure.ai.projects.aio import AIProjectClient @@ -38,6 +38,11 @@ from agent_framework_foundry._oauth_helpers import try_parse_oauth_consent_event +from ._feature_usage import ( + FeatureIndex, + create_feature_usage_user_agent_policy, + create_foundry_feature_usage_http_client, +) from ._tools import _sanitize_foundry_response_tool # pyright: ignore[reportPrivateUsage] if sys.version_info >= (3, 13): @@ -174,6 +179,7 @@ class MyClient(FunctionInvocationLayer, RawFoundryAgentChatClient): """ OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.AGENT def __init__( self, @@ -251,8 +257,10 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "user_agent": get_user_agent(), + "user_agent_policy": create_feature_usage_user_agent_policy(), } + if IS_TELEMETRY_ENABLED: + project_client_kwargs["user_agent"] = get_user_agent() if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview self.project_client = AIProjectClient(**project_client_kwargs) @@ -261,6 +269,7 @@ def __init__( openai_client_kwargs: dict[str, Any] = {} if default_headers: openai_client_kwargs["default_headers"] = dict(default_headers) + openai_client_kwargs["http_client"] = create_foundry_feature_usage_http_client() if allow_preview: openai_client_kwargs["agent_name"] = self.agent_name openai_client = self.project_client.get_openai_client(**openai_client_kwargs) @@ -799,7 +808,9 @@ async def create_conversation(self, *, session_id: str | None = None) -> AgentSe Foundry conversation ID. """ client = cast(RawFoundryAgentChatClient, self.client) - conversation = await client.project_client.get_openai_client().conversations.create() + conversation = await client.project_client.get_openai_client( + http_client=create_foundry_feature_usage_http_client() + ).conversations.create() return self.get_session(service_session_id=conversation.id, session_id=session_id) @override diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index b55d2d101e4..8f909f004a6 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -17,7 +17,7 @@ ) from agent_framework._compaction import CompactionStrategy, TokenizerProtocol from agent_framework._feature_stage import ExperimentalFeature, experimental -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import IS_TELEMETRY_ENABLED, get_user_agent from agent_framework.observability import ChatTelemetryLayer from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient from azure.ai.projects.aio import AIProjectClient @@ -56,6 +56,11 @@ from agent_framework_foundry._oauth_helpers import try_parse_oauth_consent_event +from ._feature_usage import ( + FeatureIndex, + create_feature_usage_user_agent_policy, + create_foundry_feature_usage_http_client, +) from ._tools import _sanitize_foundry_response_tool # pyright: ignore[reportPrivateUsage] if sys.version_info >= (3, 13): @@ -151,6 +156,7 @@ class RawFoundryChatClient( OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = False + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.CHAT_CLIENT def __init__( self, @@ -220,8 +226,10 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": project_endpoint, "credential": credential, - "user_agent": get_user_agent(), + "user_agent_policy": create_feature_usage_user_agent_policy(), } + if IS_TELEMETRY_ENABLED: + project_client_kwargs["user_agent"] = get_user_agent() if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview project_client = AIProjectClient(**project_client_kwargs) @@ -229,6 +237,7 @@ def __init__( openai_kwargs: dict[str, Any] = {} if default_headers: openai_kwargs["default_headers"] = default_headers + openai_kwargs["http_client"] = create_foundry_feature_usage_http_client() super().__init__( model=resolved_model, diff --git a/python/packages/foundry/agent_framework_foundry/_embedding_client.py b/python/packages/foundry/agent_framework_foundry/_embedding_client.py index cc9668ec4fa..7e5ed4b4a0c 100644 --- a/python/packages/foundry/agent_framework_foundry/_embedding_client.py +++ b/python/packages/foundry/agent_framework_foundry/_embedding_client.py @@ -17,11 +17,14 @@ UsageDetails, load_settings, ) +from agent_framework._telemetry import IS_TELEMETRY_ENABLED, get_user_agent, mark_feature_used from agent_framework.observability import EmbeddingTelemetryLayer from azure.ai.inference.aio import EmbeddingsClient, ImageEmbeddingsClient from azure.ai.inference.models import ImageEmbeddingInput from azure.core.credentials import AzureKeyCredential +from ._feature_usage import FeatureIndex, create_feature_usage_user_agent_policy + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -151,14 +154,15 @@ def __init__( if credential is None and text_client is None and image_client is None: raise ValueError("Either 'api_key', 'credential', or pre-configured client(s) must be provided.") - self._text_client = text_client or EmbeddingsClient( - endpoint=resolved_endpoint, # type: ignore[arg-type] - credential=credential, # type: ignore[arg-type] - ) - self._image_client = image_client or ImageEmbeddingsClient( - endpoint=resolved_endpoint, # type: ignore[arg-type] - credential=credential, # type: ignore[arg-type] - ) + client_kwargs: dict[str, Any] = { + "endpoint": resolved_endpoint, + "credential": credential, + "user_agent_policy": create_feature_usage_user_agent_policy(), + } + if IS_TELEMETRY_ENABLED: + client_kwargs["user_agent"] = get_user_agent() + self._text_client = text_client or EmbeddingsClient(**client_kwargs) + self._image_client = image_client or ImageEmbeddingsClient(**client_kwargs) self._endpoint = resolved_endpoint super().__init__(additional_properties=additional_properties) @@ -206,6 +210,7 @@ async def get_embeddings( """ if not values: return GeneratedEmbeddings([], options=options) + mark_feature_used(FeatureIndex.EMBEDDING) opts: dict[str, Any] = dict(options) if options else {} diff --git a/python/packages/foundry/agent_framework_foundry/_feature_usage.py b/python/packages/foundry/agent_framework_foundry/_feature_usage.py new file mode 100644 index 00000000000..260a924abfe --- /dev/null +++ b/python/packages/foundry/agent_framework_foundry/_feature_usage.py @@ -0,0 +1,63 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum +from typing import Any + +from agent_framework._telemetry import ( + IS_TELEMETRY_ENABLED, + USER_AGENT_KEY, + apply_feature_token, + get_user_agent, + remove_feature_token, +) +from agent_framework_openai._feature_usage import ( + _is_approved_origin, # pyright: ignore[reportPrivateUsage] + create_feature_usage_http_client, +) +from azure.core.pipeline.policies import UserAgentPolicy +from openai import DefaultAsyncHttpxClient + + +class FeatureIndex(IntEnum): + """Foundry-owned feature-usage indexes.""" + + CHAT_CLIENT = 48 + AGENT = 49 + MEMORY = 50 + EMBEDDING = 51 + EVALS = 52 + TOOLBOX = 53 + + +_FOUNDRY_ORIGIN_SUFFIXES = ( + "inference.ai.azure.com", + "services.ai.azure.com", +) + + +def create_foundry_feature_usage_http_client() -> DefaultAsyncHttpxClient: + """Create an OpenAI SDK client for approved Foundry origins.""" + return create_feature_usage_http_client(approved_origin_suffixes=_FOUNDRY_ORIGIN_SUFFIXES) + + +def create_feature_usage_user_agent_policy() -> "FeatureUsageUserAgentPolicy": + """Create the Azure policy with the Agent Framework base User-Agent when enabled.""" + if IS_TELEMETRY_ENABLED: + return FeatureUsageUserAgentPolicy(user_agent=get_user_agent()) + return FeatureUsageUserAgentPolicy() + + +class FeatureUsageUserAgentPolicy(UserAgentPolicy): + """Refresh or remove the feature token based on the actual Azure request origin.""" + + def on_request(self, request: Any) -> None: + """Apply normal Azure User-Agent behavior, then destination-aware feature stamping.""" + super().on_request(request) + headers = request.http_request.headers + user_agent = headers.get(USER_AGENT_KEY) + base_user_agent = user_agent if isinstance(user_agent, str) else get_user_agent() + headers[USER_AGENT_KEY] = ( + apply_feature_token(base_user_agent) + if _is_approved_origin(request.http_request.url, _FOUNDRY_ORIGIN_SUFFIXES) + else remove_feature_token(base_user_agent) + ) diff --git a/python/packages/foundry/agent_framework_foundry/_foundry_evals.py b/python/packages/foundry/agent_framework_foundry/_foundry_evals.py index 09fa1316fe3..6fa7cad15ed 100644 --- a/python/packages/foundry/agent_framework_foundry/_foundry_evals.py +++ b/python/packages/foundry/agent_framework_foundry/_foundry_evals.py @@ -43,9 +43,11 @@ RubricScore, ) from agent_framework._feature_stage import ExperimentalFeature, experimental +from agent_framework._telemetry import mark_feature_used from openai import AsyncOpenAI from ._chat_client import FoundryChatClient +from ._feature_usage import FeatureIndex, create_foundry_feature_usage_http_client if TYPE_CHECKING: from azure.ai.projects.aio import AIProjectClient @@ -665,7 +667,7 @@ def _resolve_openai_client( return client.client return client if project_client is not None: - oai = project_client.get_openai_client() + oai = project_client.get_openai_client(http_client=create_foundry_feature_usage_http_client()) if oai is None: # pyright: ignore[reportUnnecessaryComparison] raise ValueError("project_client.get_openai_client() returned None. Check project configuration.") if not isinstance(oai, AsyncOpenAI): @@ -716,7 +718,14 @@ async def _evaluate_via_responses_impl( data_source=data_source, # type: ignore[arg-type] ) - return await _poll_eval_run(client, eval_obj.id, run.id, poll_interval, timeout, provider=provider) + return await _poll_eval_run( + client, + eval_obj.id, + run.id, + poll_interval, + timeout, + provider=provider, + ) # --------------------------------------------------------------------------- @@ -879,6 +888,7 @@ async def evaluate( Returns: ``EvalResults`` with status, counts, and portal link. """ + mark_feature_used(FeatureIndex.EVALS) # Resolve evaluators with auto-detection resolved = _resolve_default_evaluators(self._evaluators, items=items) # Filter tool evaluators if items don't have tools @@ -1023,6 +1033,7 @@ async def evaluate_traces( ) """ oai_client = _resolve_openai_client(client, project_client) + mark_feature_used(FeatureIndex.EVALS) resolved_evaluators = _resolve_default_evaluators(evaluators) if response_ids: @@ -1060,7 +1071,13 @@ async def evaluate_traces( data_source=trace_source, # type: ignore[arg-type] ) - return await _poll_eval_run(oai_client, eval_obj.id, run.id, poll_interval, timeout) + return await _poll_eval_run( + oai_client, + eval_obj.id, + run.id, + poll_interval, + timeout, + ) @experimental(feature_id=ExperimentalFeature.EVALS) @@ -1109,6 +1126,7 @@ async def evaluate_foundry_target( if "type" not in target: raise ValueError("target dict must include a 'type' key (e.g., 'azure_ai_agent').") oai_client = _resolve_openai_client(client, project_client) + mark_feature_used(FeatureIndex.EVALS) resolved_evaluators = _resolve_default_evaluators(evaluators) eval_obj = await oai_client.evals.create( @@ -1135,4 +1153,10 @@ async def evaluate_foundry_target( data_source=data_source, # type: ignore[arg-type] ) - return await _poll_eval_run(oai_client, eval_obj.id, run.id, poll_interval, timeout) + return await _poll_eval_run( + oai_client, + eval_obj.id, + run.id, + poll_interval, + timeout, + ) diff --git a/python/packages/foundry/agent_framework_foundry/_memory_provider.py b/python/packages/foundry/agent_framework_foundry/_memory_provider.py index 0397b792947..b8241369be5 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -20,12 +20,14 @@ SessionContext, load_settings, ) -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import IS_TELEMETRY_ENABLED, get_user_agent, mark_feature_used from azure.ai.projects.aio import AIProjectClient from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential from openai.types.responses import ResponseInputItemParam +from ._feature_usage import FeatureIndex, create_feature_usage_user_agent_policy + if sys.version_info >= (3, 11): from typing import Self, TypedDict # pragma: no cover else: @@ -119,8 +121,10 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "user_agent": get_user_agent(), + "user_agent_policy": create_feature_usage_user_agent_policy(), } + if IS_TELEMETRY_ENABLED: + project_client_kwargs["user_agent"] = get_user_agent() if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview project_client = AIProjectClient(**project_client_kwargs) @@ -164,6 +168,7 @@ async def before_run( 2. Searches for contextual memories based on input messages 3. Combines and injects memories into the context """ + mark_feature_used(FeatureIndex.MEMORY) # On first run, retrieve static memories (user profile memories) if not state.get("initialized"): try: diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index f837108de15..4ab555cb3d6 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable from types import SimpleNamespace from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch from uuid import uuid4 import httpx @@ -31,6 +31,7 @@ tool, ) from agent_framework_openai._chat_client import RawOpenAIChatClient +from agent_framework_openai._feature_usage import FeatureIndex as OpenAIFeatureIndex from azure.ai.projects import models as projects_models from azure.core.exceptions import ResourceNotFoundError from azure.identity import AzureCliCredential @@ -44,6 +45,7 @@ _FoundryAgentChatClient, ) from agent_framework_foundry._chat_client import FoundryChatClient +from agent_framework_foundry._feature_usage import FeatureIndex skip_if_foundry_agent_integration_tests_disabled = pytest.mark.skipif( os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/") @@ -60,6 +62,11 @@ ) +def test_raw_foundry_agent_chat_client_does_not_mark_openai_feature() -> None: + assert RawOpenAIChatClient._FEATURE_USAGE_INDEX is OpenAIFeatureIndex.OPENAI + assert RawFoundryAgentChatClient._FEATURE_USAGE_INDEX is FeatureIndex.AGENT + + def _get_foundry_azure_ai_search_model() -> str | None: """Return the model/deployment to use for local Azure AI Search integration validation.""" return next((os.environ[key] for key in _FOUNDRY_AZURE_AI_SEARCH_MODEL_ENV_VARS if os.getenv(key)), None) @@ -116,7 +123,7 @@ def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None: assert client.agent_name == "test-agent" assert client.agent_version == "1.0" - mock_project.get_openai_client.assert_called_once_with() + mock_project.get_openai_client.assert_called_once_with(http_client=ANY) async def test_foundry_agent_basic_call_does_not_request_unsupported_encrypted_reasoning() -> None: @@ -214,6 +221,7 @@ def test_raw_foundry_agent_chat_client_init_passes_agent_name_when_preview_enabl mock_project.get_openai_client.assert_called_once_with( agent_name="hosted-agent", default_headers={"x-test": "1"}, + http_client=ANY, ) diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index d8a8babfb93..6d08386137a 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -11,9 +11,10 @@ from typing import Annotated, Any, cast from unittest.mock import AsyncMock, MagicMock, patch +import agent_framework._telemetry as telemetry import pytest from agent_framework import Agent, ChatResponse, Content, Message, SupportsChatGetResponse, tool -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework.exceptions import ChatClientException, ChatClientInvalidRequestException from agent_framework_openai import OpenAIContentFilterException from agent_framework_openai._chat_client import RawOpenAIChatClient @@ -25,6 +26,7 @@ from pytest import param from agent_framework_foundry import FoundryChatClient, RawFoundryChatClient +from agent_framework_foundry._feature_usage import FeatureIndex, FeatureUsageUserAgentPolicy class OutputStruct(BaseModel): @@ -34,6 +36,24 @@ class OutputStruct(BaseModel): weather: str | None = None +def test_foundry_feature_usage_policy_refreshes_user_agent() -> None: + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + mark_feature_used(FeatureIndex.CHAT_CLIENT) + request = MagicMock() + request.http_request.url = "https://project.services.ai.azure.com/api/projects/test" + request.http_request.headers = {"User-Agent": "azsdk-python-ai-projects/1.0 agent-framework-python/1.0"} + FeatureUsageUserAgentPolicy().on_request(request) + + assert request.http_request.headers["User-Agent"] == ( + "azsdk-python-ai-projects/1.0 agent-framework-python/1.0 (feat=v1.1000000000000)" + ) + + +def test_raw_foundry_chat_client_owns_foundry_feature_bit() -> None: + assert RawFoundryChatClient._FEATURE_USAGE_INDEX is FeatureIndex.CHAT_CLIENT + + @tool(approval_mode="never_require") async def get_weather(location: Annotated[str, "The location as a city name"]) -> str: """Get the current weather in a given location.""" diff --git a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py index f005a737dc9..433abee5cfd 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py @@ -5,10 +5,11 @@ import os from collections.abc import Sequence from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest from agent_framework import Content +from agent_framework._telemetry import get_user_agent from agent_framework_foundry import ( FoundryEmbeddingClient, @@ -200,12 +201,24 @@ def test_settings_from_env(self) -> None: }, clear=True, ), - patch("agent_framework_foundry._embedding_client.EmbeddingsClient"), - patch("agent_framework_foundry._embedding_client.ImageEmbeddingsClient"), + patch("agent_framework_foundry._embedding_client.EmbeddingsClient") as text_client_type, + patch("agent_framework_foundry._embedding_client.ImageEmbeddingsClient") as image_client_type, ): client = RawFoundryEmbeddingClient() assert client.model == "env-model" assert client.image_model == "env-model" # falls back to model + text_client_type.assert_called_once_with( + endpoint="https://env.inference.ai.azure.com", + credential=ANY, + user_agent=get_user_agent(), + user_agent_policy=ANY, + ) + image_client_type.assert_called_once_with( + endpoint="https://env.inference.ai.azure.com", + credential=ANY, + user_agent=get_user_agent(), + user_agent_policy=ANY, + ) def test_image_model_from_env(self) -> None: """image_model is loaded from its own environment variable.""" diff --git a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py index 89d9023602f..6084161d94a 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py +++ b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py @@ -5,7 +5,7 @@ import os from typing import Any, cast -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import ANY, AsyncMock, Mock, patch import pytest from agent_framework import AgentResponse, Message @@ -97,6 +97,7 @@ def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMoc credential=mock_credential, allow_preview=True, user_agent=get_user_agent(), + user_agent_policy=ANY, ) diff --git a/python/packages/foundry_local/agent_framework_foundry_local/_feature_usage.py b/python/packages/foundry_local/agent_framework_foundry_local/_feature_usage.py new file mode 100644 index 00000000000..78dda4d7df9 --- /dev/null +++ b/python/packages/foundry_local/agent_framework_foundry_local/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Foundry Local-owned feature-usage indexes.""" + + FOUNDRY_LOCAL = 54 diff --git a/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py b/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py index 2d627a51649..a36647b8e35 100644 --- a/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py +++ b/python/packages/foundry_local/agent_framework_foundry_local/_foundry_local_client.py @@ -4,7 +4,7 @@ import sys from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import Any, Generic, Literal, cast, overload +from typing import Any, ClassVar, Generic, Literal, cast, overload from agent_framework import ( ChatAndFunctionMiddlewareTypes, @@ -27,6 +27,8 @@ from openai import AsyncOpenAI from pydantic import BaseModel +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -139,6 +141,8 @@ class FoundryLocalClient( ): """Foundry Local Chat completion class with middleware, telemetry, and function invocation support.""" + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.FOUNDRY_LOCAL + @overload def get_response( self, diff --git a/python/packages/foundry_local/tests/test_foundry_local_client.py b/python/packages/foundry_local/tests/test_foundry_local_client.py index e610562646a..2ba5e9d89c0 100644 --- a/python/packages/foundry_local/tests/test_foundry_local_client.py +++ b/python/packages/foundry_local/tests/test_foundry_local_client.py @@ -9,11 +9,16 @@ from agent_framework.exceptions import SettingNotFoundError from agent_framework.foundry import FoundryLocalClient +from agent_framework_foundry_local._feature_usage import FeatureIndex from agent_framework_foundry_local._foundry_local_client import FoundryLocalSettings # Settings Tests +def test_foundry_local_owns_foundry_local_feature() -> None: + assert FoundryLocalClient._FEATURE_USAGE_INDEX is FeatureIndex.FOUNDRY_LOCAL + + def test_foundry_local_settings_init_from_env(foundry_local_unit_test_env: dict[str, str]) -> None: """Test FoundryLocalSettings initialization from environment variables.""" settings = load_settings(FoundryLocalSettings, env_prefix="FOUNDRY_LOCAL_") diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index fce5ecdb3ba..14efdcdca93 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -38,7 +38,7 @@ ) from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer from agent_framework._settings import SecretString -from agent_framework._telemetry import USER_AGENT_KEY +from agent_framework._telemetry import USER_AGENT_KEY, mark_feature_used from agent_framework._tools import ( SHELL_TOOL_KIND_VALUE, FunctionInvocationConfiguration, @@ -91,6 +91,7 @@ from pydantic import BaseModel from ._exceptions import OpenAIContentFilterException +from ._feature_usage import FeatureIndex from ._shared import ( AzureTokenProvider, _attach_prompt_cache_breakpoint, # pyright: ignore[reportPrivateUsage] @@ -394,6 +395,7 @@ class RawOpenAIChatClient( INJECTABLE: ClassVar[set[str]] = {"client"} STORES_BY_DEFAULT: ClassVar[bool] = True SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = True + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.OPENAI # Azure OpenAI Responses API may include this header in responses naming the actual model that # served the request (e.g. ``gpt-5-nano-2025-08-07``), which can differ from the deployment alias @@ -624,6 +626,8 @@ async def _prepare_request( Tuple of (client, run_options, validated_options). """ client = self.client + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) validated_options = await self._validate_options(options) run_options = await self._prepare_options(messages, validated_options) return client, run_options, validated_options @@ -653,6 +657,7 @@ def _inner_get_response( **kwargs: Any, ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: continuation_token: OpenAIContinuationToken | None = options.get("continuation_token") + extra_headers = cast("Mapping[str, Any] | None", kwargs.get("extra_headers")) if stream: function_call_ids: dict[int, tuple[str, str]] = {} @@ -673,12 +678,15 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: if continuation_token is not None: # Resume a background streaming response by retrieving with stream=True client = self.client + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) validated_options = await self._validate_options(options) response_format = validated_options.get("response_format") try: raw_stream_response = await client.responses.with_raw_response.retrieve( continuation_token["response_id"], stream=True, + extra_headers=extra_headers, ) # Read headers defensively: telemetry instrumentors (e.g. azure-ai-projects # experimental tracing) wrap the streaming response in objects that do not @@ -704,6 +712,8 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: run_options, validated_options, ) = await self._prepare_request(messages, options) + if extra_headers is not None: + run_options["extra_headers"] = dict(extra_headers) response_format = validated_options.get("response_format") try: if "text_format" in run_options: @@ -748,9 +758,14 @@ async def _get_response() -> ChatResponse: if continuation_token is not None: # Poll a background response by retrieving without stream client = self.client + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) validated_options = await self._validate_options(options) try: - raw_response = await client.responses.with_raw_response.retrieve(continuation_token["response_id"]) + raw_response = await client.responses.with_raw_response.retrieve( + continuation_token["response_id"], + extra_headers=extra_headers, + ) response = raw_response.parse() except Exception as ex: self._handle_request_error(ex) @@ -769,6 +784,8 @@ async def _get_response() -> ChatResponse: options.pop("continuation_token", None) return chat_response client, run_options, validated_options = await self._prepare_request(messages, options) + if extra_headers is not None: + run_options["extra_headers"] = dict(extra_headers) try: if "text_format" in run_options: raw_response = await client.responses.with_raw_response.parse(stream=False, **run_options) diff --git a/python/packages/openai/agent_framework_openai/_chat_completion_client.py b/python/packages/openai/agent_framework_openai/_chat_completion_client.py index fd025341fc1..a28b34e1ce0 100644 --- a/python/packages/openai/agent_framework_openai/_chat_completion_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_completion_client.py @@ -23,7 +23,7 @@ from agent_framework._docstrings import apply_layered_docstring from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer from agent_framework._settings import SecretString -from agent_framework._telemetry import USER_AGENT_KEY +from agent_framework._telemetry import USER_AGENT_KEY, mark_feature_used from agent_framework._tools import ( FunctionInvocationConfiguration, FunctionInvocationLayer, @@ -59,6 +59,7 @@ from pydantic import BaseModel from ._exceptions import OpenAIContentFilterException +from ._feature_usage import FeatureIndex from ._shared import ( PROMPT_CACHE_BREAKPOINT_KEY, AzureTokenProvider, @@ -238,6 +239,7 @@ class RawOpenAIChatCompletionClient( """ INJECTABLE: ClassVar[set[str]] = {"client"} + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.OPENAI @overload def __init__( @@ -566,15 +568,20 @@ def _inner_get_response( ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: # prepare options_dict = self._prepare_options(messages, options) + extra_headers = cast("Mapping[str, Any] | None", kwargs.get("extra_headers")) if stream: - # Streaming mode - options_dict["stream_options"] = {"include_usage": True} async def _stream() -> AsyncIterable[ChatResponseUpdate]: client = self.client + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) + request_options = dict(options_dict) + request_options["stream_options"] = {"include_usage": True} + if extra_headers is not None: + request_options["extra_headers"] = dict(extra_headers) try: - async for chunk in await client.chat.completions.create(stream=True, **options_dict): + async for chunk in await client.chat.completions.create(stream=True, **request_options): if len(chunk.choices) == 0 and chunk.usage is None: continue yield self._parse_response_update_from_openai(chunk) @@ -605,9 +612,14 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # Non-streaming mode async def _get_response() -> ChatResponse: client = self.client + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) + request_options = dict(options_dict) + if extra_headers is not None: + request_options["extra_headers"] = dict(extra_headers) try: return self._parse_response_from_openai( - await client.chat.completions.create(stream=False, **options_dict), options + await client.chat.completions.create(stream=False, **request_options), options ) except BadRequestError as ex: if ex.code == "content_filter": diff --git a/python/packages/openai/agent_framework_openai/_embedding_client.py b/python/packages/openai/agent_framework_openai/_embedding_client.py index b847eb0fd15..572e2437cb1 100644 --- a/python/packages/openai/agent_framework_openai/_embedding_client.py +++ b/python/packages/openai/agent_framework_openai/_embedding_client.py @@ -10,11 +10,12 @@ from agent_framework._clients import BaseEmbeddingClient from agent_framework._settings import SecretString -from agent_framework._telemetry import USER_AGENT_KEY +from agent_framework._telemetry import USER_AGENT_KEY, mark_feature_used from agent_framework._types import Embedding, EmbeddingGenerationOptions, GeneratedEmbeddings, UsageDetails from agent_framework.observability import EmbeddingTelemetryLayer from openai import AsyncAzureOpenAI, AsyncOpenAI +from ._feature_usage import FeatureIndex from ._shared import AzureTokenProvider, load_openai_service_settings if sys.version_info >= (3, 13): @@ -68,6 +69,7 @@ class RawOpenAIEmbeddingClient( """Raw OpenAI embedding client without telemetry.""" INJECTABLE: ClassVar[set[str]] = {"client"} + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.OPENAI @overload def __init__( @@ -276,6 +278,8 @@ async def get_embeddings( raise ValueError("model is required") kwargs: dict[str, Any] = {"input": list(values), "model": model} + if self._FEATURE_USAGE_INDEX is not None: + mark_feature_used(self._FEATURE_USAGE_INDEX) if dimensions := opts.get("dimensions"): kwargs["dimensions"] = dimensions if encoding_format := opts.get("encoding_format"): diff --git a/python/packages/openai/agent_framework_openai/_feature_usage.py b/python/packages/openai/agent_framework_openai/_feature_usage.py new file mode 100644 index 00000000000..59208b0b96e --- /dev/null +++ b/python/packages/openai/agent_framework_openai/_feature_usage.py @@ -0,0 +1,44 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + +import httpx +from agent_framework._telemetry import USER_AGENT_KEY, apply_feature_token, remove_feature_token +from openai import DefaultAsyncHttpxClient + + +class FeatureIndex(IntEnum): + """OpenAI-owned feature-usage indexes.""" + + OPENAI = 56 + + +_AZURE_OPENAI_ORIGIN_SUFFIXES = ( + "cognitiveservices.azure.com", + "openai.azure.com", + "services.ai.azure.com", +) + + +def _is_approved_origin(url: httpx.URL | str, suffixes: tuple[str, ...]) -> bool: + if isinstance(url, str): + url = httpx.URL(url) + host = (url.host or "").rstrip(".").lower() + return url.scheme == "https" and any(host == suffix or host.endswith(f".{suffix}") for suffix in suffixes) + + +def create_feature_usage_http_client( + *, + approved_origin_suffixes: tuple[str, ...] = _AZURE_OPENAI_ORIGIN_SUFFIXES, +) -> DefaultAsyncHttpxClient: + """Create the OpenAI SDK default client with destination-aware feature stamping.""" + + async def stamp_feature_usage(request: httpx.Request) -> None: # ruff:ignore[unused-async] + user_agent = request.headers.get(USER_AGENT_KEY, "") + request.headers[USER_AGENT_KEY] = ( + apply_feature_token(user_agent) + if _is_approved_origin(request.url, approved_origin_suffixes) + else remove_feature_token(user_agent) + ) + + return DefaultAsyncHttpxClient(event_hooks={"request": [stamp_feature_usage]}) diff --git a/python/packages/openai/agent_framework_openai/_shared.py b/python/packages/openai/agent_framework_openai/_shared.py index 80f2ec2b598..b2ee84aacfc 100644 --- a/python/packages/openai/agent_framework_openai/_shared.py +++ b/python/packages/openai/agent_framework_openai/_shared.py @@ -18,6 +18,8 @@ from openai.types.responses.response import Response from openai.types.responses.response_stream_event import ResponseStreamEvent +from ._feature_usage import create_feature_usage_http_client + if sys.version_info >= (3, 11): from typing import TypedDict # pragma: no cover else: @@ -285,6 +287,7 @@ def load_openai_service_settings( if client: return azure_settings, client, True # type: ignore[return-value] client_args["default_headers"] = merged_headers + client_args["http_client"] = create_feature_usage_http_client() if endpoint := azure_settings.get("endpoint"): if responses_mode: client_args["base_url"] = f"{endpoint.rstrip('/')}/openai/v1/" @@ -317,6 +320,7 @@ def load_openai_service_settings( openai_args: dict[str, Any] = { "base_url": resolved_base_url, "default_headers": client_args.get("default_headers"), + "http_client": client_args["http_client"], } if "azure_ad_token_provider" in client_args: openai_args["api_key"] = _ensure_async_token_provider(client_args["azure_ad_token_provider"]) diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py index c6b85ca5ab6..f769b46716e 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py @@ -7,12 +7,14 @@ from typing import Any, cast from unittest.mock import MagicMock, patch +import agent_framework._telemetry as telemetry import pytest from agent_framework import ( Agent, ChatResponse, Content, Message, + ResponseStream, SupportsChatGetResponse, SupportsCodeInterpreterTool, SupportsFileSearchTool, @@ -21,6 +23,8 @@ SupportsWebSearchTool, tool, ) +from agent_framework._telemetry import FeatureIndex as CoreFeatureIndex +from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import ChatClientException, SettingNotFoundError from openai import BadRequestError from openai.types.chat.chat_completion import ChatCompletion, Choice @@ -33,6 +37,7 @@ _AZURE_WEB_SEARCH_UNSUPPORTED_MSG, ) from agent_framework_openai._exceptions import OpenAIContentFilterException +from agent_framework_openai._feature_usage import FeatureIndex skip_if_openai_integration_tests_disabled = pytest.mark.skipif( os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"), @@ -1895,6 +1900,38 @@ async def test_streaming_exception_handling( pass +async def test_streaming_feature_is_marked_when_request_is_sent( + openai_unit_test_env: dict[str, str], +) -> None: + client = OpenAIChatCompletionClient() + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + + async def create(**kwargs: Any) -> Any: + async def chunks() -> Any: + if False: + yield None + + return chunks() + + with patch.object(client.client.chat.completions, "create", side_effect=create): + stream = client._inner_get_response( + messages=[Message(role="user", contents=["test"])], + stream=True, + options={}, + ) + assert isinstance(stream, ResponseStream) + mark_feature_used(CoreFeatureIndex.CORE_AGENT) + async for _ in stream: + pass + + token = telemetry.get_feature_token() + assert token is not None + mask = int(token.split(".", 1)[1], 16) + assert mask & (1 << FeatureIndex.OPENAI) + assert mask & (1 << CoreFeatureIndex.CORE_AGENT) + + # region Integration Tests diff --git a/python/packages/openai/tests/openai/test_openai_embedding_client.py b/python/packages/openai/tests/openai/test_openai_embedding_client.py index 8347f9fd3f1..0325101e0e2 100644 --- a/python/packages/openai/tests/openai/test_openai_embedding_client.py +++ b/python/packages/openai/tests/openai/test_openai_embedding_client.py @@ -7,8 +7,10 @@ from typing import Any, cast from unittest.mock import AsyncMock, MagicMock +import agent_framework._telemetry as telemetry import pytest from agent_framework import SupportsGetEmbeddings +from agent_framework._telemetry import get_feature_token from agent_framework.exceptions import SettingNotFoundError from openai.types import CreateEmbeddingResponse from openai.types import Embedding as OpenAIEmbedding @@ -19,6 +21,7 @@ OpenAIEmbeddingOptions, ) from agent_framework_openai._embedding_client import RawOpenAIEmbeddingClient +from agent_framework_openai._feature_usage import FeatureIndex def _make_openai_response( @@ -106,6 +109,24 @@ async def test_openai_get_embeddings(openai_unit_test_env: dict[str, str]) -> No assert result[0].dimensions == 3 +async def test_embedding_request_marks_openai_feature( + openai_unit_test_env: dict[str, str], +) -> None: + mock_response = _make_openai_response(embeddings=[[0.1]]) + client = OpenAIEmbeddingClient() + client.client = MagicMock() + client.client.embeddings = MagicMock() + client.client.embeddings.create = AsyncMock(return_value=mock_response) + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + + await client.get_embeddings(["test"]) + + token = get_feature_token() + assert token is not None + assert int(token.split(".", 1)[1], 16) & (1 << FeatureIndex.OPENAI) + + async def test_openai_get_embeddings_usage(openai_unit_test_env: dict[str, str]) -> None: mock_response = _make_openai_response( embeddings=[[0.1]], diff --git a/python/packages/openai/tests/openai/test_openai_shared.py b/python/packages/openai/tests/openai/test_openai_shared.py index c3e17e8ac8b..e485c55a62b 100644 --- a/python/packages/openai/tests/openai/test_openai_shared.py +++ b/python/packages/openai/tests/openai/test_openai_shared.py @@ -6,10 +6,16 @@ from typing import Any, cast from unittest.mock import MagicMock, patch +import agent_framework._telemetry as telemetry +import httpx import pytest +from agent_framework import AGENT_FRAMEWORK_USER_AGENT +from agent_framework._telemetry import FeatureIndex as CoreFeatureIndex +from agent_framework._telemetry import mark_feature_used from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential +from agent_framework_openai._feature_usage import create_feature_usage_http_client from agent_framework_openai._shared import ( AZURE_OPENAI_TOKEN_SCOPE, _ensure_async_token_provider, @@ -71,6 +77,31 @@ def test_resolve_azure_invalid_credential_raises() -> None: _resolve_azure_credential_to_token_provider(cast(Any, object())) +async def test_feature_usage_hook_stamps_approved_origin_and_strips_custom_origin() -> None: + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + mark_feature_used(CoreFeatureIndex.CORE_AGENT) + client = create_feature_usage_http_client() + hook = client.event_hooks["request"][0] + approved = httpx.Request( + "POST", + "https://resource.openai.azure.com/openai/v1/responses", + headers={"User-Agent": f"{AGENT_FRAMEWORK_USER_AGENT} sdk/1.0"}, + ) + custom = httpx.Request( + "POST", + "https://customer-gateway.example.com/v1/responses", + headers={"User-Agent": f"{AGENT_FRAMEWORK_USER_AGENT} (feat=v1.1)"}, + ) + + await hook(approved) + await hook(custom) + await client.aclose() + + assert approved.headers["User-Agent"] == f"{AGENT_FRAMEWORK_USER_AGENT} sdk/1.0 (feat=v1.1)" + assert custom.headers["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT + + async def test_ensure_async_token_provider_wraps_sync_provider() -> None: def sync_provider() -> str: return "sync-token" diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py b/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py index a1c299bf723..9c5c8b2ab05 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py @@ -7,16 +7,17 @@ from typing import Any, Literal, cast from agent_framework import AgentResponse, Message, SupportsAgentRun +from agent_framework._telemetry import mark_feature_used from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse from agent_framework._workflows._agent_utils import resolve_agent_id from agent_framework._workflows._checkpoint import CheckpointStorage from agent_framework._workflows._executor import Executor, handler from agent_framework._workflows._message_utils import normalize_messages_input from agent_framework._workflows._workflow import Workflow -from agent_framework._workflows._workflow_builder import WorkflowBuilder from agent_framework._workflows._workflow_context import WorkflowContext from typing_extensions import Never +from ._feature_usage import FeatureIndex from ._orchestration_request_info import AgentApprovalExecutor from ._participant_output_config import ( UNSET, @@ -26,6 +27,7 @@ _ParticipantOutputSpecifier, # pyright: ignore[reportPrivateUsage] _resolve_participant_output_config, # pyright: ignore[reportPrivateUsage] ) +from ._workflow_builder import OrchestrationWorkflowBuilder as WorkflowBuilder logger = logging.getLogger(__name__) @@ -402,6 +404,7 @@ def build(self) -> Workflow: workflow = ConcurrentBuilder(participants=[agent1, agent2]).build() """ + mark_feature_used(FeatureIndex.CONCURRENT) # Internal nodes dispatcher = _DispatchToAllParticipants(id="dispatcher") aggregator = self._aggregator if self._aggregator is not None else _AggregateAgentConversations(id="aggregator") diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py b/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py new file mode 100644 index 00000000000..86d804a061b --- /dev/null +++ b/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Orchestration-owned feature-usage indexes.""" + + SEQUENTIAL = 32 + CONCURRENT = 33 + GROUP_CHAT = 34 + MAGENTIC = 35 + HANDOFF = 36 diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py index 3e5a2a97758..d3cac2158fd 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py @@ -30,12 +30,12 @@ from typing import Any, ClassVar, Literal, cast from agent_framework import Agent, AgentResponse, AgentResponseUpdate, AgentSession, Message, SupportsAgentRun +from agent_framework._telemetry import mark_feature_used from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse from agent_framework._workflows._agent_utils import resolve_agent_id from agent_framework._workflows._checkpoint import CheckpointStorage from agent_framework._workflows._executor import Executor from agent_framework._workflows._workflow import Workflow -from agent_framework._workflows._workflow_builder import WorkflowBuilder from agent_framework._workflows._workflow_context import WorkflowContext from pydantic import BaseModel, Field from typing_extensions import Never @@ -49,6 +49,7 @@ ParticipantRegistry, TerminationCondition, ) +from ._feature_usage import FeatureIndex from ._orchestration_request_info import AgentApprovalExecutor from ._orchestrator_helpers import clean_conversation_for_handoff from ._participant_output_config import ( @@ -59,6 +60,7 @@ _ParticipantOutputSpecifier, # pyright: ignore[reportPrivateUsage] _resolve_participant_output_config, # pyright: ignore[reportPrivateUsage] ) +from ._workflow_builder import OrchestrationWorkflowBuilder as WorkflowBuilder if sys.version_info >= (3, 12): from typing import override # pragma: no cover @@ -1010,6 +1012,7 @@ def build(self) -> Workflow: Returns: Validated Workflow instance ready for execution """ + mark_feature_used(FeatureIndex.GROUP_CHAT) # Resolve orchestrator and participants to executors participants: list[Executor] = self._resolve_participants() orchestrator: Executor = self._resolve_orchestrator(participants) diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py b/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py index 4e6eb571c1b..04d40432c72 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py @@ -41,6 +41,7 @@ from agent_framework import Agent, AgentResponse, Message, SupportsAgentRun from agent_framework._middleware import FunctionInvocationContext, FunctionMiddleware, MiddlewareTermination from agent_framework._sessions import AgentSession +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import FunctionTool, tool from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest from agent_framework._workflows._agent_utils import resolve_agent_id @@ -48,10 +49,10 @@ from agent_framework._workflows._events import WorkflowEvent from agent_framework._workflows._request_info_mixin import response_handler from agent_framework._workflows._workflow import Workflow -from agent_framework._workflows._workflow_builder import WorkflowBuilder from agent_framework._workflows._workflow_context import WorkflowContext from ._base_group_chat_orchestrator import TerminationCondition +from ._feature_usage import FeatureIndex from ._orchestrator_helpers import clean_conversation_for_handoff from ._participant_output_config import ( UNSET, @@ -61,6 +62,7 @@ _ParticipantOutputSpecifier, # pyright: ignore[reportPrivateUsage] _resolve_participant_output_config, # pyright: ignore[reportPrivateUsage] ) +from ._workflow_builder import OrchestrationWorkflowBuilder as WorkflowBuilder if sys.version_info >= (3, 12): from typing import override # pragma: no cover @@ -941,6 +943,7 @@ def build(self) -> Workflow: ValueError: If participants or coordinator were not configured, or if required configuration is invalid. """ + mark_feature_used(FeatureIndex.HANDOFF) # Resolve agents (either from instances or factories) # The returned map keys are either executor IDs or factory names, which is need to resolve handoff configs resolved_agents = self._resolve_agents() diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py index 14ae940002e..902b26d652c 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py @@ -18,6 +18,7 @@ Message, SupportsAgentRun, ) +from agent_framework._telemetry import mark_feature_used from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse from agent_framework._workflows._checkpoint import CheckpointStorage from agent_framework._workflows._events import WorkflowEvent @@ -25,7 +26,6 @@ from agent_framework._workflows._model_utils import DictConvertible, encode_value from agent_framework._workflows._request_info_mixin import response_handler from agent_framework._workflows._workflow import Workflow -from agent_framework._workflows._workflow_builder import WorkflowBuilder from agent_framework._workflows._workflow_context import WorkflowContext from typing_extensions import Never, Sentinel @@ -37,6 +37,7 @@ GroupChatWorkflowContextOutT, ParticipantRegistry, ) +from ._feature_usage import FeatureIndex from ._participant_output_config import ( UNSET, _coalesce_output_from, # pyright: ignore[reportPrivateUsage] @@ -45,6 +46,7 @@ _ParticipantOutputSpecifier, # pyright: ignore[reportPrivateUsage] _resolve_participant_output_config, # pyright: ignore[reportPrivateUsage] ) +from ._workflow_builder import OrchestrationWorkflowBuilder as WorkflowBuilder if sys.version_info >= (3, 12): from typing import override # pragma: no cover @@ -1774,6 +1776,7 @@ def _resolve_participants(self) -> list[Executor]: def build(self) -> Workflow: """Build a Magentic workflow with the orchestrator and all agent executors.""" + mark_feature_used(FeatureIndex.MAGENTIC) logger.info(f"Building Magentic workflow with {len(self._participants)} participants") participants: list[Executor] = self._resolve_participants() diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_orchestration_request_info.py b/python/packages/orchestrations/agent_framework_orchestrations/_orchestration_request_info.py index 66949ae5769..234e64b284e 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_orchestration_request_info.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_orchestration_request_info.py @@ -10,10 +10,11 @@ from agent_framework._workflows._executor import Executor, handler from agent_framework._workflows._request_info_mixin import response_handler from agent_framework._workflows._workflow import Workflow -from agent_framework._workflows._workflow_builder import WorkflowBuilder from agent_framework._workflows._workflow_context import WorkflowContext from agent_framework._workflows._workflow_executor import WorkflowExecutor +from ._workflow_builder import OrchestrationWorkflowBuilder as WorkflowBuilder + def resolve_request_info_filter(agents: list[str | SupportsAgentRun] | None) -> set[str]: """Resolve a list of agent/executor references to a set of IDs for filtering. diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py b/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py index 4f8720b0bf3..4da4a3a7c20 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py @@ -19,6 +19,7 @@ from typing import Any, Literal, cast from agent_framework import Message, SupportsAgentRun +from agent_framework._telemetry import mark_feature_used from agent_framework._workflows._agent_executor import AgentExecutor from agent_framework._workflows._agent_utils import resolve_agent_id from agent_framework._workflows._checkpoint import CheckpointStorage @@ -28,9 +29,9 @@ ) from agent_framework._workflows._message_utils import normalize_messages_input from agent_framework._workflows._workflow import Workflow -from agent_framework._workflows._workflow_builder import WorkflowBuilder from agent_framework._workflows._workflow_context import WorkflowContext +from ._feature_usage import FeatureIndex from ._orchestration_request_info import AgentApprovalExecutor from ._participant_output_config import ( UNSET, @@ -40,6 +41,7 @@ _ParticipantOutputSpecifier, # pyright: ignore[reportPrivateUsage] _resolve_participant_output_config, # pyright: ignore[reportPrivateUsage] ) +from ._workflow_builder import OrchestrationWorkflowBuilder as WorkflowBuilder logger = logging.getLogger(__name__) @@ -241,6 +243,7 @@ def build(self) -> Workflow: terminator's own `yield_output` is Workflow Output (`AgentResponse`, or per-chunk `AgentResponseUpdate` when streaming). """ + mark_feature_used(FeatureIndex.SEQUENTIAL) input_conv = _InputToConversation(id="input-conversation") # Resolve participants and participant factories to executors diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_workflow_builder.py b/python/packages/orchestrations/agent_framework_orchestrations/_workflow_builder.py new file mode 100644 index 00000000000..5313eb3bc5e --- /dev/null +++ b/python/packages/orchestrations/agent_framework_orchestrations/_workflow_builder.py @@ -0,0 +1,12 @@ +# Copyright (c) Microsoft. All rights reserved. + +from typing import ClassVar + +from agent_framework._telemetry import FeatureIndex as CoreFeatureIndex +from agent_framework._workflows._workflow_builder import WorkflowBuilder + + +class OrchestrationWorkflowBuilder(WorkflowBuilder): + """Workflow builder that leaves usage attribution to the orchestration.""" + + _FEATURE_USAGE_INDEX: ClassVar[CoreFeatureIndex | None] = None diff --git a/python/packages/orchestrations/tests/test_magentic.py b/python/packages/orchestrations/tests/test_magentic.py index 6551f72495d..ec586a1ca3a 100644 --- a/python/packages/orchestrations/tests/test_magentic.py +++ b/python/packages/orchestrations/tests/test_magentic.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from typing import Any, ClassVar, cast +import agent_framework._telemetry as telemetry import pytest from agent_framework import ( Agent, @@ -27,6 +28,7 @@ WorkflowRunState, handler, ) +from agent_framework._telemetry import get_feature_token from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage from agent_framework.orchestrations import ( GroupChatRequestMessage, @@ -40,6 +42,8 @@ StandardMagenticManager, ) +from agent_framework_orchestrations._feature_usage import FeatureIndex + if sys.version_info >= (3, 12): from typing import override # type: ignore # pragma: no cover else: @@ -186,6 +190,17 @@ async def _noop( pass +def test_magentic_builder_marks_feature_with_custom_manager() -> None: + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + + MagenticBuilder(participants=[DummyExec("agentA")], manager=FakeManager()).build() + + token = get_feature_token() + assert token is not None + assert int(token.split(".", 1)[1], 16) & (1 << FeatureIndex.MAGENTIC) + + async def test_magentic_builder_returns_workflow_and_runs() -> None: manager = FakeManager() agent = StubAgent(manager.next_speaker_name, "first draft") diff --git a/python/packages/orchestrations/tests/test_sequential.py b/python/packages/orchestrations/tests/test_sequential.py index 73b35b554fa..758762ea89f 100644 --- a/python/packages/orchestrations/tests/test_sequential.py +++ b/python/packages/orchestrations/tests/test_sequential.py @@ -3,6 +3,7 @@ from collections.abc import AsyncIterable, Awaitable, Sequence from typing import Any, Literal, overload +import agent_framework._telemetry as telemetry import pytest from agent_framework import ( AgentExecutorResponse, @@ -20,9 +21,13 @@ WorkflowRunState, handler, ) +from agent_framework._telemetry import FeatureIndex as CoreFeatureIndex +from agent_framework._telemetry import get_feature_token from agent_framework._workflows._checkpoint import InMemoryCheckpointStorage from agent_framework.orchestrations import SequentialBuilder +from agent_framework_orchestrations._feature_usage import FeatureIndex + class _EchoAgent(BaseAgent): """Simple agent that appends a single assistant message with its name.""" @@ -91,6 +96,19 @@ async def summarize(self, conversation: list[str], ctx: WorkflowContext[list[Mes pass +def test_sequential_builder_does_not_mark_custom_workflow() -> None: + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + + SequentialBuilder(participants=[_EchoAgent(name="echo")]).build() + + token = get_feature_token() + assert token is not None + mask = int(token.split(".", 1)[1], 16) + assert mask & (1 << FeatureIndex.SEQUENTIAL) + assert not mask & (1 << CoreFeatureIndex.CORE_WORKFLOW) + + def test_sequential_builder_rejects_empty_participants() -> None: with pytest.raises(ValueError): SequentialBuilder(participants=[]) diff --git a/python/samples/README.md b/python/samples/README.md index 4a585db7813..19324c85e97 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -123,6 +123,8 @@ variable. | `agent-framework-core` | `observability` | `ENABLE_INSTRUMENTATION` | `true` | | `agent-framework-core` | `observability` | `ENABLE_SENSITIVE_DATA` | `false` | | `agent-framework-core` | `observability` | `ENABLE_CONSOLE_EXPORTERS` | `true` | +| `agent-framework-core` | `agent_framework._telemetry` | `AGENT_FRAMEWORK_FEATURE_MASK_DISABLED` | `true` | +| `agent-framework-core` | `agent_framework._telemetry` | `AGENT_FRAMEWORK_USER_AGENT_DISABLED` | `true` | | `agent-framework-core` | `observability` | `OTEL_EXPORTER_OTLP_ENDPOINT` | `http://localhost:4317` | | `agent-framework-core` | `observability` | `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` | `http://localhost:4318/v1/traces` | | `agent-framework-core` | `observability` | `OTEL_EXPORTER_OTLP_METRICS_ENDPOINT` | `http://localhost:4318/v1/metrics` | From c5fc32a694e38f1fca79423c49763daed9462fde Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Fri, 24 Jul 2026 13:07:20 +0200 Subject: [PATCH 2/8] Python: track declarative feature usage Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- .../_feature_usage.py | 10 ++++++ .../agent_framework_declarative/_loader.py | 10 ++++-- .../_workflows/_factory.py | 3 ++ .../tests/test_declarative_loader.py | 32 +++++++++++++++++++ .../tests/test_workflow_factory.py | 20 ++++++++++++ 5 files changed, 73 insertions(+), 2 deletions(-) create mode 100644 python/packages/declarative/agent_framework_declarative/_feature_usage.py diff --git a/python/packages/declarative/agent_framework_declarative/_feature_usage.py b/python/packages/declarative/agent_framework_declarative/_feature_usage.py new file mode 100644 index 00000000000..631164b1822 --- /dev/null +++ b/python/packages/declarative/agent_framework_declarative/_feature_usage.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Declarative-owned feature-usage indexes.""" + + AGENT = 75 + WORKFLOW = 76 diff --git a/python/packages/declarative/agent_framework_declarative/_loader.py b/python/packages/declarative/agent_framework_declarative/_loader.py index bbd9f4ec9be..dc3bff2aa26 100644 --- a/python/packages/declarative/agent_framework_declarative/_loader.py +++ b/python/packages/declarative/agent_framework_declarative/_loader.py @@ -19,9 +19,11 @@ ExperimentalFeature, experimental, ) +from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentException from dotenv import load_dotenv +from ._feature_usage import FeatureIndex from ._models import ( AnonymousConnection, ApiKeyConnection, @@ -471,13 +473,15 @@ def create_agent_from_dict(self, agent_def: dict[str, Any]) -> Agent: if output_schema := prompt_agent.outputSchema: chat_options["response_format"] = output_schema.to_json_schema() # Step 3: Create the agent instance - return Agent( + agent = Agent( client=client, name=prompt_agent.name, description=prompt_agent.description, instructions=prompt_agent.instructions, default_options=chat_options, # type: ignore[arg-type] ) + mark_feature_used(FeatureIndex.AGENT) + return agent async def create_agent_from_yaml_path_async(self, yaml_path: str | Path) -> Agent: """Async version: Create a Agent from a YAML file path. @@ -582,13 +586,15 @@ async def create_agent_from_dict_async(self, agent_def: dict[str, Any]) -> Agent chat_options["tools"] = tools if output_schema := prompt_agent.outputSchema: chat_options["response_format"] = output_schema.to_json_schema() - return Agent( + agent = Agent( client=client, name=prompt_agent.name, description=prompt_agent.description, instructions=prompt_agent.instructions, default_options=chat_options, # type: ignore[arg-type] ) + mark_feature_used(FeatureIndex.AGENT) + return agent async def _create_agent_with_provider(self, prompt_agent: PromptAgent, mapping: ProviderTypeMapping) -> Agent: """Create an Agent through a provider object that exposes ``create_agent``. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index 10206b13c54..09e847f193c 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -24,7 +24,9 @@ SupportsAgentRun, Workflow, ) +from agent_framework._telemetry import mark_feature_used +from .._feature_usage import FeatureIndex from .._loader import AgentFactory from ._declarative_base import DeclarativeEnvConfig, discover_env_references from ._declarative_builder import DeclarativeWorkflowBuilder @@ -471,6 +473,7 @@ def _create_workflow( len(graph_builder._executors), # type: ignore[reportPrivateUsage] ) + mark_feature_used(FeatureIndex.WORKFLOW) return workflow def _normalize_workflow_def(self, workflow_def: dict[str, Any]) -> dict[str, Any]: diff --git a/python/packages/declarative/tests/test_declarative_loader.py b/python/packages/declarative/tests/test_declarative_loader.py index 823342f68ea..8e63974cc8e 100644 --- a/python/packages/declarative/tests/test_declarative_loader.py +++ b/python/packages/declarative/tests/test_declarative_loader.py @@ -492,6 +492,38 @@ def test_create_agent_from_dict_parses_prompt_agent(self): assert agent is not None + def test_create_agent_from_dict_marks_declarative_agent_used(self): + """Test that successful declarative agent creation marks feature usage.""" + from agent_framework_declarative import AgentFactory + from agent_framework_declarative._feature_usage import FeatureIndex + + factory = AgentFactory(client=MagicMock()) + + with patch("agent_framework_declarative._loader.mark_feature_used") as mark_feature_used: + factory.create_agent_from_dict({ + "kind": "Prompt", + "name": "TestAgent", + "instructions": "You are a helpful assistant.", + }) + + mark_feature_used.assert_called_once_with(FeatureIndex.AGENT) + + async def test_create_agent_from_dict_async_marks_declarative_agent_used(self): + """Test that successful async declarative agent creation marks feature usage.""" + from agent_framework_declarative import AgentFactory + from agent_framework_declarative._feature_usage import FeatureIndex + + factory = AgentFactory(client=MagicMock()) + + with patch("agent_framework_declarative._loader.mark_feature_used") as mark_feature_used: + await factory.create_agent_from_dict_async({ + "kind": "Prompt", + "name": "TestAgent", + "instructions": "You are a helpful assistant.", + }) + + mark_feature_used.assert_called_once_with(FeatureIndex.AGENT) + def test_create_agent_from_dict_matches_yaml(self): """Test that create_agent_from_dict produces same result as create_agent_from_yaml.""" from unittest.mock import MagicMock diff --git a/python/packages/declarative/tests/test_workflow_factory.py b/python/packages/declarative/tests/test_workflow_factory.py index acf677f6c84..2306e223e4e 100644 --- a/python/packages/declarative/tests/test_workflow_factory.py +++ b/python/packages/declarative/tests/test_workflow_factory.py @@ -3,9 +3,11 @@ """Unit tests for WorkflowFactory.""" from typing import Any, cast +from unittest.mock import patch import pytest +from agent_framework_declarative._feature_usage import FeatureIndex from agent_framework_declarative._workflows._errors import DeclarativeWorkflowError from agent_framework_declarative._workflows._factory import WorkflowFactory @@ -66,6 +68,24 @@ def test_valid_minimal_workflow(self): assert workflow is not None assert workflow.name == "minimal-workflow" + def test_valid_workflow_marks_declarative_workflow_used(self): + """Test that successful declarative workflow creation marks feature usage.""" + factory = WorkflowFactory() + + with patch("agent_framework_declarative._workflows._factory.mark_feature_used") as mark_feature_used: + factory.create_workflow_from_definition({ + "name": "minimal-workflow", + "actions": [ + { + "kind": "SetValue", + "path": "Local.result", + "value": "done", + } + ], + }) + + mark_feature_used.assert_called_once_with(FeatureIndex.WORKFLOW) + @_requires_powerfx class TestWorkflowFactoryExecution: From 8c9ceafbd88bc7ba08545477f5612e9a1549b24f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 29 Jul 2026 16:50:27 +0200 Subject: [PATCH 3/8] Python: complete feature usage telemetry Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- .../a2a/agent_framework_a2a/_a2a_executor.py | 3 + .../a2a/agent_framework_a2a/_agent.py | 4 ++ .../a2a/agent_framework_a2a/_feature_usage.py | 9 +++ .../ag-ui/agent_framework_ag_ui/_agent.py | 3 + .../ag-ui/agent_framework_ag_ui/_client.py | 3 + .../agent_framework_ag_ui/_feature_usage.py | 9 +++ .../ag-ui/agent_framework_ag_ui/_workflow.py | 3 + .../agent_framework_anthropic/_chat_client.py | 6 +- .../_feature_usage.py | 9 +++ .../anthropic/tests/test_anthropic_client.py | 9 ++- .../_context_provider.py | 5 +- .../_feature_usage.py | 9 +++ .../tests/test_aisearch_context_provider.py | 12 ++++ .../_context_provider.py | 4 ++ .../_feature_usage.py | 9 +++ .../tests/cu/test_context_provider.py | 10 ++- .../_context_provider.py | 7 +++ .../_feature_usage.py | 9 +++ .../tests/test_context_provider.py | 17 +++++ .../_checkpoint_storage.py | 6 +- .../_feature_usage.py | 9 +++ .../_history_provider.py | 6 +- .../tests/test_cosmos_history_provider.py | 10 +++ .../agent_framework_azurefunctions/_app.py | 3 + .../_feature_usage.py | 9 +++ .../agent_framework_bedrock/_chat_client.py | 5 +- .../_embedding_client.py | 5 +- .../agent_framework_bedrock/_feature_usage.py | 9 +++ .../bedrock/tests/test_bedrock_client.py | 5 +- .../agent_framework_chatkit/_converter.py | 4 ++ .../agent_framework_chatkit/_feature_usage.py | 9 +++ .../agent_framework_chatkit/_streaming.py | 4 ++ .../claude/agent_framework_claude/_agent.py | 4 ++ .../agent_framework_claude/_feature_usage.py | 9 +++ .../claude/tests/test_claude_agent.py | 7 ++- .../agent_framework_copilotstudio/_agent.py | 4 ++ .../_feature_usage.py | 9 +++ .../copilotstudio/tests/test_copilot_agent.py | 5 +- .../core/agent_framework/_telemetry.py | 2 +- .../core/tests/core/test_telemetry.py | 60 ++++++++++++++---- .../_feature_usage.py | 4 +- .../agent_framework_declarative/_loader.py | 4 +- .../_workflows/_factory.py | 2 +- .../tests/test_declarative_loader.py | 4 +- .../tests/test_workflow_factory.py | 2 +- .../devui/agent_framework_devui/__init__.py | 4 ++ .../agent_framework_devui/_feature_usage.py | 9 +++ .../_feature_usage.py | 9 +++ .../agent_framework_durabletask/_shim.py | 3 + .../agent_framework_durabletask/_worker.py | 3 + .../foundry/agent_framework_foundry/_agent.py | 6 +- .../agent_framework_foundry/_chat_client.py | 6 +- .../_embedding_client.py | 6 +- .../agent_framework_foundry/_feature_usage.py | 35 +++++------ .../agent_framework_foundry/_foundry_evals.py | 6 +- .../_memory_provider.py | 6 +- .../tests/foundry/test_foundry_agent.py | 19 +++++- .../tests/foundry/test_foundry_chat_client.py | 63 +++++++++++++++++-- .../foundry/test_foundry_embedding_client.py | 4 +- .../foundry/test_foundry_memory_provider.py | 2 +- .../_feature_usage.py | 10 +++ .../_invocations.py | 4 ++ .../_responses.py | 4 ++ .../_toolbox.py | 9 +++ .../foundry_hosting/tests/test_toolbox.py | 26 +++++++- .../agent_framework_gemini/_chat_client.py | 6 +- .../agent_framework_gemini/_feature_usage.py | 9 +++ .../gemini/tests/test_gemini_client.py | 5 +- .../agent_framework_github_copilot/_agent.py | 5 ++ .../_feature_usage.py | 9 +++ .../tests/test_github_copilot_agent.py | 5 +- .../_conversion.py | 5 ++ .../_feature_usage.py | 9 +++ .../_agent_tool.py | 4 ++ .../_feature_usage.py | 9 +++ .../_workflow_tool.py | 4 ++ .../_feature_usage.py | 9 +++ .../_parsing.py | 6 ++ .../_feature_usage.py | 9 +++ .../_parsing.py | 4 ++ .../_rendering.py | 5 ++ .../agent_framework_hosting/_feature_usage.py | 9 +++ .../hosting/agent_framework_hosting/_state.py | 5 ++ .../_feature_usage.py | 9 +++ .../agent_framework_hyperlight/_provider.py | 3 + .../agent_framework_lab_common/__init__.py | 1 + .../_feature_usage.py | 9 +++ .../lab/gaia/agent_framework_lab_gaia/gaia.py | 5 ++ .../agent_framework_lab_lightning/__init__.py | 3 + python/packages/lab/pyproject.toml | 9 ++- .../tau2/agent_framework_lab_tau2/runner.py | 3 + .../agent_framework_mem0/_context_provider.py | 5 ++ .../agent_framework_mem0/_feature_usage.py | 9 +++ .../mem0/tests/test_mem0_context_provider.py | 9 ++- .../_embedding_client.py | 4 ++ .../agent_framework_mistral/_feature_usage.py | 9 +++ .../mistral/test_mistral_embedding_client.py | 7 ++- .../agent_framework_monty/_feature_usage.py | 9 +++ .../monty/agent_framework_monty/_provider.py | 3 + .../agent_framework_ollama/_chat_client.py | 5 ++ .../_embedding_client.py | 4 ++ .../agent_framework_ollama/_feature_usage.py | 9 +++ .../ollama/test_ollama_embedding_client.py | 7 ++- .../_concurrent.py | 2 +- .../_feature_usage.py | 10 +-- .../_group_chat.py | 2 +- .../_handoff.py | 2 +- .../_magentic.py | 2 +- .../_sequential.py | 2 +- .../orchestrations/tests/test_magentic.py | 2 +- .../orchestrations/tests/test_sequential.py | 2 +- .../agent_framework_purview/_client.py | 7 ++- .../agent_framework_purview/_feature_usage.py | 9 +++ .../tests/purview/test_purview_client.py | 7 ++- .../_context_provider.py | 5 ++ .../agent_framework_redis/_feature_usage.py | 9 +++ .../_history_provider.py | 5 ++ python/packages/redis/tests/test_providers.py | 10 +++ .../agent_framework_tools/_feature_usage.py | 9 +++ .../agent_framework_tools/shell/_docker.py | 3 + .../agent_framework_tools/shell/_tool.py | 3 + 121 files changed, 799 insertions(+), 105 deletions(-) create mode 100644 python/packages/a2a/agent_framework_a2a/_feature_usage.py create mode 100644 python/packages/ag-ui/agent_framework_ag_ui/_feature_usage.py create mode 100644 python/packages/anthropic/agent_framework_anthropic/_feature_usage.py create mode 100644 python/packages/azure-ai-search/agent_framework_azure_ai_search/_feature_usage.py create mode 100644 python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_feature_usage.py create mode 100644 python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_feature_usage.py create mode 100644 python/packages/azure-cosmos/agent_framework_azure_cosmos/_feature_usage.py create mode 100644 python/packages/azurefunctions/agent_framework_azurefunctions/_feature_usage.py create mode 100644 python/packages/bedrock/agent_framework_bedrock/_feature_usage.py create mode 100644 python/packages/chatkit/agent_framework_chatkit/_feature_usage.py create mode 100644 python/packages/claude/agent_framework_claude/_feature_usage.py create mode 100644 python/packages/copilotstudio/agent_framework_copilotstudio/_feature_usage.py create mode 100644 python/packages/devui/agent_framework_devui/_feature_usage.py create mode 100644 python/packages/durabletask/agent_framework_durabletask/_feature_usage.py create mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_feature_usage.py create mode 100644 python/packages/gemini/agent_framework_gemini/_feature_usage.py create mode 100644 python/packages/github_copilot/agent_framework_github_copilot/_feature_usage.py create mode 100644 python/packages/hosting-a2a/agent_framework_hosting_a2a/_feature_usage.py create mode 100644 python/packages/hosting-mcp/agent_framework_hosting_mcp/_feature_usage.py create mode 100644 python/packages/hosting-responses/agent_framework_hosting_responses/_feature_usage.py create mode 100644 python/packages/hosting-telegram/agent_framework_hosting_telegram/_feature_usage.py create mode 100644 python/packages/hosting/agent_framework_hosting/_feature_usage.py create mode 100644 python/packages/hyperlight/agent_framework_hyperlight/_feature_usage.py create mode 100644 python/packages/lab/common/agent_framework_lab_common/__init__.py create mode 100644 python/packages/lab/common/agent_framework_lab_common/_feature_usage.py create mode 100644 python/packages/mem0/agent_framework_mem0/_feature_usage.py create mode 100644 python/packages/mistral/agent_framework_mistral/_feature_usage.py create mode 100644 python/packages/monty/agent_framework_monty/_feature_usage.py create mode 100644 python/packages/ollama/agent_framework_ollama/_feature_usage.py create mode 100644 python/packages/purview/agent_framework_purview/_feature_usage.py create mode 100644 python/packages/redis/agent_framework_redis/_feature_usage.py create mode 100644 python/packages/tools/agent_framework_tools/_feature_usage.py diff --git a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py index 55a7b620cdd..2aa028d3856 100644 --- a/python/packages/a2a/agent_framework_a2a/_a2a_executor.py +++ b/python/packages/a2a/agent_framework_a2a/_a2a_executor.py @@ -19,8 +19,10 @@ Message, SupportsAgentRun, ) +from agent_framework._telemetry import mark_feature_used from typing_extensions import override +from ._feature_usage import FeatureIndex from ._utils import get_uri_data logger = logging.getLogger("agent_framework.a2a") @@ -147,6 +149,7 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non if context.message is None: raise ValueError("Message must be provided in the RequestContext") + mark_feature_used(FeatureIndex.A2A) query = context.get_user_input() task = context.current_task diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index 4beeb10749c..19966f94069 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -42,10 +42,12 @@ normalize_messages, prepend_agent_framework_to_user_agent, ) +from agent_framework._telemetry import mark_feature_used from agent_framework._types import AgentRunInputs from agent_framework.observability import AgentTelemetryLayer from google.protobuf.json_format import MessageToDict +from ._feature_usage import FeatureIndex from ._utils import get_uri_data if sys.version_info >= (3, 11): @@ -542,6 +544,8 @@ async def _map_a2a_stream( session: The agent session for context providers. session_context: The session context for context providers. """ + mark_feature_used(FeatureIndex.A2A) + if session_context is None: session_context = SessionContext(input_messages=[], options={}) diff --git a/python/packages/a2a/agent_framework_a2a/_feature_usage.py b/python/packages/a2a/agent_framework_a2a/_feature_usage.py new file mode 100644 index 00000000000..2eb6d2bc2f7 --- /dev/null +++ b/python/packages/a2a/agent_framework_a2a/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """A2A-owned feature-usage indexes.""" + + A2A = 71 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 f67e04a420c..09766b7e1aa 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -7,9 +7,11 @@ from ag_ui.core import BaseEvent from agent_framework import SupportsAgentRun +from agent_framework._telemetry import mark_feature_used from ._agent_run import PendingApprovalEntry, PendingApprovalKey, run_agent_stream from ._approval_state import InMemoryAGUIApprovalStateStore +from ._feature_usage import FeatureIndex from ._snapshots import AGUIThreadSnapshotStore @@ -142,6 +144,7 @@ async def run( Yields: AG-UI events """ + mark_feature_used(FeatureIndex.AG_UI) async for event in run_agent_stream( input_data, self.agent, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index 98ecc36dc25..f7edb458277 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -24,10 +24,12 @@ ResponseStream, ) from agent_framework._middleware import ChatMiddlewareLayer +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer from agent_framework.observability import ChatTelemetryLayer from ._event_converters import AGUIEventConverter +from ._feature_usage import FeatureIndex from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume from ._message_adapters import agent_framework_messages_to_agui from ._utils import convert_tools_to_agui_format @@ -398,6 +400,7 @@ async def _streaming_impl( Yields: ChatResponseUpdate objects """ + mark_feature_used(FeatureIndex.AG_UI) messages_to_send, state = self._extract_state_from_messages(messages) thread_id = self._get_thread_id(options) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_feature_usage.py b/python/packages/ag-ui/agent_framework_ag_ui/_feature_usage.py new file mode 100644 index 00000000000..07f7f7900fb --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """AG-UI-owned feature-usage indexes.""" + + AG_UI = 72 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 27b3f4fa626..20f44361c91 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -25,7 +25,9 @@ ToolCallStartEvent, ) from agent_framework import Workflow +from agent_framework._telemetry import mark_feature_used +from ._feature_usage import FeatureIndex from ._message_adapters import agui_messages_to_snapshot_format from ._run_common import ( _build_run_finished_event, @@ -298,6 +300,7 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: Subclasses may override this to provide custom AG-UI streams. """ + mark_feature_used(FeatureIndex.AG_UI) thread_id = self._thread_id_from_input(input_data) run_id = str(input_data.get("run_id") or input_data.get("runId") or uuid.uuid4()) snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY)) diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index abda651e31a..e323eaee748 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -28,7 +28,7 @@ tool, ) from agent_framework._settings import SecretString, load_settings -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework._tools import SHELL_TOOL_KIND_VALUE, normalize_tools from agent_framework._types import _get_data_bytes_as_str # type: ignore from agent_framework.observability import ChatTelemetryLayer @@ -55,6 +55,8 @@ from anthropic.types.beta.beta_encrypted_code_execution_result_block import BetaEncryptedCodeExecutionResultBlock from pydantic import BaseModel +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 11): from typing import TypedDict # pragma: no cover else: @@ -550,6 +552,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # each message_delta carries the running total), so thread a per-stream # accumulator to _process_stream_event to emit increments instead. emitted_usage: dict[str, int] = {} + mark_feature_used(FeatureIndex.ANTHROPIC) async for chunk in await self.anthropic_client.beta.messages.create(**run_options, stream=True): parsed_chunk = self._process_stream_event(chunk, emitted_usage) if parsed_chunk: @@ -559,6 +562,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: # Non-streaming mode async def _get_response() -> ChatResponse: + mark_feature_used(FeatureIndex.ANTHROPIC) message = await self.anthropic_client.beta.messages.create(**run_options, stream=False) return self._process_message(message, options) diff --git a/python/packages/anthropic/agent_framework_anthropic/_feature_usage.py b/python/packages/anthropic/agent_framework_anthropic/_feature_usage.py new file mode 100644 index 00000000000..a2c779f377c --- /dev/null +++ b/python/packages/anthropic/agent_framework_anthropic/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Anthropic-owned feature-usage indexes.""" + + ANTHROPIC = 57 diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index 3c417541408..d713111152e 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -33,6 +33,7 @@ from agent_framework_anthropic import AnthropicClient, RawAnthropicClient from agent_framework_anthropic._chat_client import AnthropicSettings +from agent_framework_anthropic._feature_usage import FeatureIndex # Test constants VALID_PNG_BASE64 = b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" @@ -1606,10 +1607,12 @@ async def test_inner_get_response(mock_anthropic_client: MagicMock) -> None: messages = [Message(role="user", contents=["Hi"])] chat_options = ChatOptions(max_tokens=10) - response = await client._inner_get_response( # type: ignore[attr-defined] - messages=messages, options=chat_options - ) + with patch("agent_framework_anthropic._chat_client.mark_feature_used") as mark_feature_used: + response = await client._inner_get_response( # type: ignore[attr-defined] + messages=messages, options=chat_options + ) + mark_feature_used.assert_called_once_with(FeatureIndex.ANTHROPIC) assert response is not None assert response.response_id == "msg_test" assert len(response.messages) == 1 diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index e93b7e061ec..d67d9d141eb 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -25,7 +25,7 @@ SupportsGetEmbeddings, load_settings, ) -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework.exceptions import SettingNotFoundError from azure.core.credentials import AzureKeyCredential, TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -47,6 +47,8 @@ VectorizedQuery, ) +from ._feature_usage import FeatureIndex + if TYPE_CHECKING: from agent_framework._agents import SupportsAgentRun from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient @@ -626,6 +628,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Retrieve relevant context from Azure AI Search and add to session context.""" + mark_feature_used(FeatureIndex.AZURE_AI_SEARCH) messages_list = list(context.input_messages) filtered_messages = [ diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_feature_usage.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_feature_usage.py new file mode 100644 index 00000000000..206239f27b0 --- /dev/null +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Azure AI Search-owned feature-usage indexes.""" + + AZURE_AI_SEARCH = 65 diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 22bb4df2a76..3aa4bf83558 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -20,6 +20,7 @@ KnowledgeBaseOutputModeLiteral, RetrievalReasoningEffortLiteral, ) +from agent_framework_azure_ai_search._feature_usage import FeatureIndex # -- Helpers ------------------------------------------------------------------- @@ -54,6 +55,17 @@ async def __anext__(self): return doc +async def test_before_run_marks_azure_ai_search_used() -> None: + provider = object.__new__(AzureAISearchContextProvider) + context = Mock(spec=SessionContext) + context.input_messages = [] + + with patch("agent_framework_azure_ai_search._context_provider.mark_feature_used") as mark_feature_used: + await provider.before_run(agent=Mock(), session=Mock(spec=AgentSession), context=context, state={}) + + mark_feature_used.assert_called_once_with(FeatureIndex.AZURE_AI_SEARCH) + + def _make_mock_index( fields: list[SimpleNamespace] | None = None, profiles: list[SimpleNamespace] | None = None, diff --git a/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_context_provider.py b/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_context_provider.py index 22b5cef5797..ce0c964cb9a 100644 --- a/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_context_provider.py +++ b/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_context_provider.py @@ -29,12 +29,15 @@ ) from agent_framework._sessions import AgentSession from agent_framework._settings import load_settings +from agent_framework._telemetry import mark_feature_used from azure.ai.contentunderstanding import to_llm_input from azure.ai.contentunderstanding.aio import ContentUnderstandingClient from azure.ai.contentunderstanding.models import AnalysisInput, AnalysisResult from azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential +from ._feature_usage import FeatureIndex + if TYPE_CHECKING: from agent_framework._agents import SupportsAgentRun @@ -275,6 +278,7 @@ async def before_run( This method is called automatically by the framework before each LLM invocation. """ + mark_feature_used(FeatureIndex.AZURE_CONTENTUNDERSTANDING) documents: dict[str, DocumentEntry] = state.setdefault("documents", {}) # Per-session mutable state — isolated per session to prevent cross-session leakage. diff --git a/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_feature_usage.py b/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_feature_usage.py new file mode 100644 index 00000000000..186167d85b9 --- /dev/null +++ b/python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Azure Content Understanding-owned feature-usage indexes.""" + + AZURE_CONTENTUNDERSTANDING = 67 diff --git a/python/packages/azure-contentunderstanding/tests/cu/test_context_provider.py b/python/packages/azure-contentunderstanding/tests/cu/test_context_provider.py index 72018408fc6..a046cfc4361 100644 --- a/python/packages/azure-contentunderstanding/tests/cu/test_context_provider.py +++ b/python/packages/azure-contentunderstanding/tests/cu/test_context_provider.py @@ -7,7 +7,7 @@ import json import re from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch from agent_framework import Content, Message, SessionContext from agent_framework._sessions import AgentSession @@ -18,6 +18,7 @@ DocumentStatus, ) from agent_framework_azure_contentunderstanding._detection import SUPPORTED_MEDIA_TYPES, derive_doc_key +from agent_framework_azure_contentunderstanding._feature_usage import FeatureIndex # --------------------------------------------------------------------------- # Helpers @@ -983,7 +984,12 @@ async def test_lazy_initialization_on_before_run(self) -> None: state: dict[str, Any] = {} session = AgentSession() - await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state) + with patch( + "agent_framework_azure_contentunderstanding._context_provider.mark_feature_used" + ) as mark_feature_used: + await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state) + + mark_feature_used.assert_called_once_with(FeatureIndex.AZURE_CONTENTUNDERSTANDING) # Client should still be set assert provider._client is not None 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 06cc1adb378..3798c3bad8b 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 @@ -17,6 +17,9 @@ from agent_framework import AgentSession, ContextProvider, Message, SessionContext from agent_framework._settings import load_settings +from agent_framework._telemetry import mark_feature_used + +from ._feature_usage import FeatureIndex if sys.version_info >= (3, 11): from typing import Self # pragma: no cover @@ -344,6 +347,8 @@ async def before_run( context: The invocation context to add memories to. state: Provider-scoped mutable state. """ + mark_feature_used(FeatureIndex.AZURE_COSMOS_MEMORY) + # Extract query from input messages query_text = "\n".join(msg.text for msg in context.input_messages if msg.text and msg.text.strip()) @@ -424,6 +429,8 @@ async def after_run( context: The invocation context with response populated. state: Provider-scoped mutable state. """ + mark_feature_used(FeatureIndex.AZURE_COSMOS_MEMORY) + # Get user_id and thread_id from provider-scoped state (falling back to the session id) user_id = self._resolve_user_id(state, session) thread_id = state.get("thread_id") or session.session_id or "default" diff --git a/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_feature_usage.py b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_feature_usage.py new file mode 100644 index 00000000000..2129647f700 --- /dev/null +++ b/python/packages/azure-cosmos-memory/agent_framework_azure_cosmos_memory/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Azure Cosmos DB memory-owned feature-usage indexes.""" + + AZURE_COSMOS_MEMORY = 82 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 f0c83ee13ba..ae682b237b9 100644 --- a/python/packages/azure-cosmos-memory/tests/test_context_provider.py +++ b/python/packages/azure-cosmos-memory/tests/test_context_provider.py @@ -23,12 +23,29 @@ DEFAULT_CONTEXT_PROMPT, CosmosMemoryContextProvider, ) +from agent_framework_azure_cosmos_memory._feature_usage import FeatureIndex # The provider methods accept an ``agent`` implementing ``SupportsAgentRun`` but never # use it in these tests, so a typed ``None`` stub keeps the call sites clean. _STUB_AGENT: Any = None +async def test_before_run_marks_cosmos_memory_used_before_empty_return() -> None: + provider = object.__new__(CosmosMemoryContextProvider) + context = MagicMock(spec=SessionContext) + context.input_messages = [] + + with patch("agent_framework_azure_cosmos_memory._context_provider.mark_feature_used") as mark_feature_used: + await provider.before_run( + agent=_STUB_AGENT, + session=MagicMock(spec=AgentSession), + context=context, + state={}, + ) + + mark_feature_used.assert_called_once_with(FeatureIndex.AZURE_COSMOS_MEMORY) + + @pytest.fixture def mock_memory_client() -> AsyncMock: """Create a mock AsyncCosmosMemoryClient.""" diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 915eee432bc..24dc9343659 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -8,7 +8,7 @@ from typing import Any, TypedDict from agent_framework._settings import SecretString, load_settings -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework._workflows._checkpoint import CheckpointID, WorkflowCheckpoint from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value from agent_framework.exceptions import WorkflowCheckpointException @@ -18,6 +18,8 @@ from azure.cosmos.aio import ContainerProxy, CosmosClient from azure.cosmos.exceptions import CosmosResourceNotFoundError +from ._feature_usage import FeatureIndex + AzureCredentialTypes = TokenCredential | AsyncTokenCredential logger = logging.getLogger(__name__) @@ -214,6 +216,7 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: Returns: The unique ID of the saved checkpoint. """ + mark_feature_used(FeatureIndex.AZURE_COSMOS) await self._ensure_container_proxy() checkpoint_dict = checkpoint.to_dict() @@ -242,6 +245,7 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: WorkflowCheckpointException: If no checkpoint with the given ID exists, or if multiple checkpoints share the same ID across workflows. """ + mark_feature_used(FeatureIndex.AZURE_COSMOS) await self._ensure_container_proxy() query = "SELECT * FROM c WHERE c.checkpoint_id = @checkpoint_id" diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_feature_usage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_feature_usage.py new file mode 100644 index 00000000000..8eff97cca1b --- /dev/null +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Azure Cosmos DB-owned feature-usage indexes.""" + + AZURE_COSMOS = 66 diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py index 62a83e6a0ff..507fe228b05 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py @@ -13,12 +13,14 @@ from agent_framework import Message from agent_framework._sessions import HistoryProvider from agent_framework._settings import SecretString, load_settings -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential from azure.cosmos import PartitionKey from azure.cosmos.aio import ContainerProxy, CosmosClient, DatabaseProxy +from ._feature_usage import FeatureIndex + AzureCredentialTypes = TokenCredential | AsyncTokenCredential logger = logging.getLogger(__name__) @@ -136,6 +138,7 @@ async def get_messages( **kwargs: Any, ) -> list[Message]: """Retrieve stored messages for this session from Azure Cosmos DB.""" + mark_feature_used(FeatureIndex.AZURE_COSMOS) await self._ensure_container_proxy() session_key = self._session_partition_key(session_id) @@ -176,6 +179,7 @@ async def save_messages( **kwargs: Any, ) -> None: """Persist messages for this session to Azure Cosmos DB.""" + mark_feature_used(FeatureIndex.AZURE_COSMOS) if not messages: return diff --git a/python/packages/azure-cosmos/tests/test_cosmos_history_provider.py b/python/packages/azure-cosmos/tests/test_cosmos_history_provider.py index bcfa7430841..7e581a7636e 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_history_provider.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_history_provider.py @@ -17,6 +17,7 @@ from azure.cosmos.exceptions import CosmosResourceNotFoundError import agent_framework_azure_cosmos._history_provider as history_provider_module +from agent_framework_azure_cosmos._feature_usage import FeatureIndex from agent_framework_azure_cosmos._history_provider import CosmosHistoryProvider skip_if_cosmos_integration_tests_disabled = pytest.mark.skipif( @@ -44,6 +45,15 @@ async def _iterator() -> AsyncIterator[Any]: return _iterator() +async def test_save_messages_marks_azure_cosmos_used_before_empty_return() -> None: + provider = object.__new__(CosmosHistoryProvider) + + with patch("agent_framework_azure_cosmos._history_provider.mark_feature_used") as mark_feature_used: + await provider.save_messages(None, []) + + mark_feature_used.assert_called_once_with(FeatureIndex.AZURE_COSMOS) + + @pytest.fixture def mock_container() -> MagicMock: container = MagicMock() diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py index 812a6788ce0..99578663c12 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py @@ -21,6 +21,7 @@ import azure.durable_functions as df import azure.functions as func from agent_framework import SupportsAgentRun, Workflow +from agent_framework._telemetry import mark_feature_used from agent_framework_durabletask import ( DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS, @@ -56,6 +57,7 @@ from ._entities import create_agent_entity from ._errors import IncomingRequestError +from ._feature_usage import FeatureIndex from ._orchestration import AgentOrchestrationContextType, AgentTask, AzureFunctionsAgentExecutor from ._routes import build_workflow_respond_url, build_workflow_status_url, split_request_url from ._workflow import run_workflow_orchestrator @@ -292,6 +294,7 @@ def __init__( if self.enable_health_check: self._setup_health_route() + mark_feature_used(FeatureIndex.AZUREFUNCTIONS) logger.debug("[AgentFunctionApp] Initialization complete") def _collect_workflows( diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_feature_usage.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_feature_usage.py new file mode 100644 index 00000000000..4bc0ffef613 --- /dev/null +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Azure Functions-owned feature-usage indexes.""" + + AZUREFUNCTIONS = 78 diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index cbe30683e28..c38b813fdae 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -31,7 +31,7 @@ validate_tool_mode, ) from agent_framework._settings import SecretString, load_settings -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework.exceptions import ChatClientInvalidResponseException from agent_framework.observability import ChatTelemetryLayer from boto3.session import Session as Boto3Session @@ -40,6 +40,8 @@ from botocore.exceptions import ClientError from pydantic import BaseModel +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -331,6 +333,7 @@ def _create_session(settings: BedrockSettings) -> Boto3Session: return Boto3Session(**session_kwargs) def _invoke_converse(self, request: Mapping[str, Any]) -> dict[str, Any]: + mark_feature_used(FeatureIndex.BEDROCK) try: response = self._bedrock_client.converse(**request) if not isinstance(response, Mapping): diff --git a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py index 4b666dbc4f2..cd71a9d0eb0 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py @@ -19,12 +19,14 @@ UsageDetails, load_settings, ) -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework.observability import EmbeddingTelemetryLayer from boto3.session import Session as Boto3Session from botocore.client import BaseClient from botocore.config import Config as BotoConfig +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -180,6 +182,7 @@ async def get_embeddings( if not model: raise ValueError("model is required") + mark_feature_used(FeatureIndex.BEDROCK) embedding_results = await asyncio.gather( *(self._generate_embedding_for_text(opts, model, text) for text in values) ) diff --git a/python/packages/bedrock/agent_framework_bedrock/_feature_usage.py b/python/packages/bedrock/agent_framework_bedrock/_feature_usage.py new file mode 100644 index 00000000000..bd8f441adc7 --- /dev/null +++ b/python/packages/bedrock/agent_framework_bedrock/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Amazon Bedrock-owned feature-usage indexes.""" + + BEDROCK = 58 diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index a2655eb877a..6d339ae3c59 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -16,6 +16,7 @@ from agent_framework_bedrock import BedrockChatClient from agent_framework_bedrock._chat_client import BedrockSettings +from agent_framework_bedrock._feature_usage import FeatureIndex class _StubBedrockRuntime: @@ -67,8 +68,10 @@ async def test_get_response_invokes_bedrock_runtime() -> None: Message(role="user", contents=[Content.from_text(text="hello")]), ] - response = await client.get_response(messages=messages, options={"max_tokens": 32}) + with patch("agent_framework_bedrock._chat_client.mark_feature_used") as mark_feature_used: + response = await client.get_response(messages=messages, options={"max_tokens": 32}) + mark_feature_used.assert_called_once_with(FeatureIndex.BEDROCK) assert stub.calls, "Expected the runtime client to be called" payload = stub.calls[0] assert payload["modelId"] == "amazon.titan-text" diff --git a/python/packages/chatkit/agent_framework_chatkit/_converter.py b/python/packages/chatkit/agent_framework_chatkit/_converter.py index c97b8479447..c4fe24670f2 100644 --- a/python/packages/chatkit/agent_framework_chatkit/_converter.py +++ b/python/packages/chatkit/agent_framework_chatkit/_converter.py @@ -11,6 +11,7 @@ Content, Message, ) +from agent_framework._telemetry import mark_feature_used from chatkit.types import ( AssistantMessageItem, Attachment, @@ -30,6 +31,8 @@ WorkflowItem, ) +from ._feature_usage import FeatureIndex + logger = logging.getLogger(__name__) @@ -610,4 +613,5 @@ async def simple_to_agent_input(thread_items: Sequence[ThreadItem] | ThreadItem) # Convert multiple items messages = await simple_to_agent_input([user_message_item, assistant_message_item, task_item]) """ + mark_feature_used(FeatureIndex.CHATKIT) return await _DEFAULT_CONVERTER.to_agent_input(thread_items) diff --git a/python/packages/chatkit/agent_framework_chatkit/_feature_usage.py b/python/packages/chatkit/agent_framework_chatkit/_feature_usage.py new file mode 100644 index 00000000000..4fdc9b02c13 --- /dev/null +++ b/python/packages/chatkit/agent_framework_chatkit/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """ChatKit-owned feature-usage indexes.""" + + CHATKIT = 73 diff --git a/python/packages/chatkit/agent_framework_chatkit/_streaming.py b/python/packages/chatkit/agent_framework_chatkit/_streaming.py index df44fa005d1..121f0af992d 100644 --- a/python/packages/chatkit/agent_framework_chatkit/_streaming.py +++ b/python/packages/chatkit/agent_framework_chatkit/_streaming.py @@ -7,6 +7,7 @@ from datetime import datetime from agent_framework import AgentResponseUpdate +from agent_framework._telemetry import mark_feature_used from chatkit.types import ( AssistantMessageContent, AssistantMessageContentPartTextDelta, @@ -17,6 +18,8 @@ ThreadStreamEvent, ) +from ._feature_usage import FeatureIndex + async def stream_agent_response( response_stream: AsyncIterable[AgentResponseUpdate], @@ -44,6 +47,7 @@ async def stream_agent_response( ThreadStreamEvent: ChatKit events representing the agent's response, including incremental text deltas for streaming display. """ + mark_feature_used(FeatureIndex.CHATKIT) # Use provided ID generator or create default one if generate_id is None: diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 7db8eb7dbfd..724c3575ce4 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -28,6 +28,7 @@ normalize_messages, normalize_tools, ) +from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentException from agent_framework.observability import AgentTelemetryLayer from claude_agent_sdk import ( @@ -42,6 +43,8 @@ ) from claude_agent_sdk.types import StreamEvent, TextBlock +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -779,6 +782,7 @@ async def _get_stream( session_id: str | None = None structured_output: Any = None + mark_feature_used(FeatureIndex.CLAUDE) await self._client.query(prompt) async for message in self._client.receive_response(): if isinstance(message, StreamEvent): diff --git a/python/packages/claude/agent_framework_claude/_feature_usage.py b/python/packages/claude/agent_framework_claude/_feature_usage.py new file mode 100644 index 00000000000..cb97adf22e8 --- /dev/null +++ b/python/packages/claude/agent_framework_claude/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Claude Agent SDK-owned feature-usage indexes.""" + + CLAUDE = 62 diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 116d4114694..4d4fe1fd039 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -9,6 +9,7 @@ from agent_framework_claude import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings from agent_framework_claude._agent import TOOLS_MCP_SERVER_NAME +from agent_framework_claude._feature_usage import FeatureIndex # region Test ClaudeAgentSettings @@ -231,9 +232,13 @@ async def test_run_with_string_message(self) -> None: ] mock_client = self._create_mock_client(messages) - with patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client): + with ( + patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client), + patch("agent_framework_claude._agent.mark_feature_used") as mark_feature_used, + ): agent = ClaudeAgent() response = await agent.run("Hello") + mark_feature_used.assert_called_once_with(FeatureIndex.CLAUDE) assert response.text == "Hello!" async def test_run_captures_session_id(self) -> None: diff --git a/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py b/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py index 0bc9c484327..1cf21740b0a 100644 --- a/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py +++ b/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py @@ -18,11 +18,13 @@ normalize_messages, ) from agent_framework._settings import load_settings +from agent_framework._telemetry import mark_feature_used from agent_framework._types import AgentRunInputs from agent_framework.exceptions import AgentException from microsoft_agents.copilotstudio.client import AgentType, ConnectionSettings, CopilotClient, PowerPlatformCloud from ._acquire_token import acquire_token +from ._feature_usage import FeatureIndex class CopilotStudioSettings(TypedDict, total=False): @@ -255,6 +257,7 @@ async def _run_impl( question = "\n".join([message.text for message in input_messages]) + mark_feature_used(FeatureIndex.COPILOTSTUDIO) activities = self.client.ask_question(question, service_session_id) response_messages: list[Message] = [] response_id: str | None = None @@ -287,6 +290,7 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: question = "\n".join([message.text for message in input_messages]) + mark_feature_used(FeatureIndex.COPILOTSTUDIO) activities = self.client.ask_question(question, service_session_id) async for message in self._process_activities(activities, streaming=True): diff --git a/python/packages/copilotstudio/agent_framework_copilotstudio/_feature_usage.py b/python/packages/copilotstudio/agent_framework_copilotstudio/_feature_usage.py new file mode 100644 index 00000000000..70249d3abf5 --- /dev/null +++ b/python/packages/copilotstudio/agent_framework_copilotstudio/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Copilot Studio-owned feature-usage indexes.""" + + COPILOTSTUDIO = 63 diff --git a/python/packages/copilotstudio/tests/test_copilot_agent.py b/python/packages/copilotstudio/tests/test_copilot_agent.py index 49da1d72086..26aba203f59 100644 --- a/python/packages/copilotstudio/tests/test_copilot_agent.py +++ b/python/packages/copilotstudio/tests/test_copilot_agent.py @@ -9,6 +9,7 @@ from microsoft_agents.copilotstudio.client import CopilotClient from agent_framework_copilotstudio import CopilotStudioAgent +from agent_framework_copilotstudio._feature_usage import FeatureIndex def create_async_generator(items: list[Any]) -> Any: @@ -136,8 +137,10 @@ async def test_run_with_string_message(self, mock_copilot_client: MagicMock, moc mock_copilot_client.start_conversation.return_value = create_async_generator([conversation_activity]) mock_copilot_client.ask_question.return_value = create_async_generator([mock_activity]) - response = await agent.run("test message") + with patch("agent_framework_copilotstudio._agent.mark_feature_used") as mark_feature_used: + response = await agent.run("test message") + mark_feature_used.assert_called_once_with(FeatureIndex.COPILOTSTUDIO) assert isinstance(response, AgentResponse) assert len(response.messages) == 1 content = response.messages[0].contents[0] diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index a7c48f3b48f..518a54c4b82 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -65,7 +65,7 @@ class FeatureIndex(IntEnum): _hosted_env_detected: bool = False _feature_mask = 0 _feature_mask_lock = threading.Lock() -_feature_comment_pattern = re.compile(r"\s+\(feat=v\d+\.[0-9a-fA-F]+\)") +_feature_comment_pattern = re.compile(r"(?:^|\s+)\(feat=v\d+\.[0-9a-fA-F]+\)") def _add_user_agent_prefix(prefix: str) -> None: diff --git a/python/packages/core/tests/core/test_telemetry.py b/python/packages/core/tests/core/test_telemetry.py index c405c930302..3add7e5c993 100644 --- a/python/packages/core/tests/core/test_telemetry.py +++ b/python/packages/core/tests/core/test_telemetry.py @@ -27,6 +27,7 @@ apply_feature_token, get_feature_token, mark_feature_used, + remove_feature_token, ) # region Test constants @@ -145,6 +146,10 @@ def test_apply_feature_token_preserves_unrelated_comments() -> None: ) +def test_remove_feature_token_strips_standalone_token() -> None: + assert remove_feature_token("(feat=v1.1)") == "" + + def test_apply_feature_token_removes_stale_token_when_disabled() -> None: _reset_feature_mask() with ( @@ -152,6 +157,7 @@ def test_apply_feature_token_removes_stale_token_when_disabled() -> None: patch.dict(os.environ, {FEATURE_MASK_DISABLED_ENV_VAR: "true"}), ): assert apply_feature_token("agent-framework-python/1.0 (feat=v1.5)") == "agent-framework-python/1.0" + assert apply_feature_token("(feat=v1.5)") == "" def test_mark_feature_used_is_thread_safe() -> None: @@ -166,7 +172,7 @@ def test_mark_feature_used_is_thread_safe() -> None: assert get_feature_token() == f"v1.{(1 << 128) - 1:x}" -def test_declared_feature_indexes_do_not_overlap() -> None: +def test_declared_feature_indexes_match_registry() -> None: registry_path = next( ( parent / "docs" / "specs" / "feature-usage-bit-registry.md" @@ -179,11 +185,20 @@ def test_declared_feature_indexes_do_not_overlap() -> None: pytest.skip("Feature-usage registry is not available outside a repository checkout.") repository_root = registry_path.parents[2] + registry_text = registry_path.read_text(encoding="utf-8") + python_table = registry_text.split("## Index table — Python", 1)[1].split("## Index table — .NET", 1)[0] + registry_rows = re.findall(r"^\| (\d+) \| `([^`]+)` \|", python_table, re.MULTILINE) + registry_pairs = {(int(index), identifier.upper().replace(".", "_")) for index, identifier in registry_rows} + assert len(registry_pairs) == len(registry_rows), "Python v1 registry contains duplicate (index, id) rows." + assert all(0 <= index < 128 for index, _ in registry_pairs) + declaration_files = [ repository_root / "python" / "packages" / "core" / "agent_framework" / "_telemetry.py", - *repository_root.glob("python/packages/*/agent_framework*/_feature_usage.py"), + *repository_root.glob("python/packages/**/_feature_usage.py"), ] - declarations: dict[int, str] = {} + declarations_by_index: dict[int, str] = {} + declaration_pairs: set[tuple[int, str]] = set() + declaration_owners: dict[tuple[int, str], tuple[Path, Path]] = {} for declaration_file in declaration_files: tree = ast.parse(declaration_file.read_text(encoding="utf-8")) for node in tree.body: @@ -198,17 +213,36 @@ def test_declared_feature_indexes_do_not_overlap() -> None: index = member.value.value if not isinstance(index, int): continue - assert 0 <= index < 128 - assert index not in declarations, ( - f"Feature index {index} overlaps between {declarations[index]} and " - f"{declaration_file.relative_to(repository_root)}:{target.id}" + declaration = f"{declaration_file.relative_to(repository_root)}:{target.id}" + assert 0 <= index < 128, f"Feature index {index} is out of range in {declaration}." + assert index not in declarations_by_index, ( + f"Feature index {index} overlaps between {declarations_by_index[index]} and {declaration}." ) - declarations[index] = f"{declaration_file.relative_to(repository_root)}:{target.id}" - - registry_text = registry_path.read_text(encoding="utf-8") - python_table = registry_text.split("## Index table — Python", 1)[1].split("## Index table — .NET", 1)[0] - registry_indexes = {int(index) for index in re.findall(r"^\| (\d+) \| `[^`]+` \|", python_table, re.MULTILINE)} - assert declarations.keys() <= registry_indexes + declarations_by_index[index] = declaration + pair = (index, target.id) + declaration_pairs.add(pair) + package_root = repository_root.joinpath(*declaration_file.relative_to(repository_root).parts[:3]) + declaration_owners[pair] = (package_root, declaration_file) + + assert declaration_pairs == registry_pairs + + for pair, (package_root, declaration_file) in declaration_owners.items(): + _, member_name = pair + referenced = False + for source_file in package_root.rglob("*.py"): + if source_file == declaration_file or "tests" in source_file.parts: + continue + source_tree = ast.parse(source_file.read_text(encoding="utf-8")) + if any( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "FeatureIndex" + and node.attr == member_name + for node in ast.walk(source_tree) + ): + referenced = True + break + assert referenced, f"Feature index {pair} is declared but never referenced by its owning package." def test_app_info_when_telemetry_enabled(): diff --git a/python/packages/declarative/agent_framework_declarative/_feature_usage.py b/python/packages/declarative/agent_framework_declarative/_feature_usage.py index 631164b1822..2d712e95ce3 100644 --- a/python/packages/declarative/agent_framework_declarative/_feature_usage.py +++ b/python/packages/declarative/agent_framework_declarative/_feature_usage.py @@ -6,5 +6,5 @@ class FeatureIndex(IntEnum): """Declarative-owned feature-usage indexes.""" - AGENT = 75 - WORKFLOW = 76 + DECLARATIVE_AGENT = 75 + DECLARATIVE_WORKFLOW = 76 diff --git a/python/packages/declarative/agent_framework_declarative/_loader.py b/python/packages/declarative/agent_framework_declarative/_loader.py index dc3bff2aa26..5be02a9ad8a 100644 --- a/python/packages/declarative/agent_framework_declarative/_loader.py +++ b/python/packages/declarative/agent_framework_declarative/_loader.py @@ -480,7 +480,7 @@ def create_agent_from_dict(self, agent_def: dict[str, Any]) -> Agent: instructions=prompt_agent.instructions, default_options=chat_options, # type: ignore[arg-type] ) - mark_feature_used(FeatureIndex.AGENT) + mark_feature_used(FeatureIndex.DECLARATIVE_AGENT) return agent async def create_agent_from_yaml_path_async(self, yaml_path: str | Path) -> Agent: @@ -593,7 +593,7 @@ async def create_agent_from_dict_async(self, agent_def: dict[str, Any]) -> Agent instructions=prompt_agent.instructions, default_options=chat_options, # type: ignore[arg-type] ) - mark_feature_used(FeatureIndex.AGENT) + mark_feature_used(FeatureIndex.DECLARATIVE_AGENT) return agent async def _create_agent_with_provider(self, prompt_agent: PromptAgent, mapping: ProviderTypeMapping) -> Agent: diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index 09e847f193c..363b4d77f37 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -473,7 +473,7 @@ def _create_workflow( len(graph_builder._executors), # type: ignore[reportPrivateUsage] ) - mark_feature_used(FeatureIndex.WORKFLOW) + mark_feature_used(FeatureIndex.DECLARATIVE_WORKFLOW) return workflow def _normalize_workflow_def(self, workflow_def: dict[str, Any]) -> dict[str, Any]: diff --git a/python/packages/declarative/tests/test_declarative_loader.py b/python/packages/declarative/tests/test_declarative_loader.py index 8e63974cc8e..cac6d4f53e5 100644 --- a/python/packages/declarative/tests/test_declarative_loader.py +++ b/python/packages/declarative/tests/test_declarative_loader.py @@ -506,7 +506,7 @@ def test_create_agent_from_dict_marks_declarative_agent_used(self): "instructions": "You are a helpful assistant.", }) - mark_feature_used.assert_called_once_with(FeatureIndex.AGENT) + mark_feature_used.assert_called_once_with(FeatureIndex.DECLARATIVE_AGENT) async def test_create_agent_from_dict_async_marks_declarative_agent_used(self): """Test that successful async declarative agent creation marks feature usage.""" @@ -522,7 +522,7 @@ async def test_create_agent_from_dict_async_marks_declarative_agent_used(self): "instructions": "You are a helpful assistant.", }) - mark_feature_used.assert_called_once_with(FeatureIndex.AGENT) + mark_feature_used.assert_called_once_with(FeatureIndex.DECLARATIVE_AGENT) def test_create_agent_from_dict_matches_yaml(self): """Test that create_agent_from_dict produces same result as create_agent_from_yaml.""" diff --git a/python/packages/declarative/tests/test_workflow_factory.py b/python/packages/declarative/tests/test_workflow_factory.py index 2306e223e4e..c40633dae1b 100644 --- a/python/packages/declarative/tests/test_workflow_factory.py +++ b/python/packages/declarative/tests/test_workflow_factory.py @@ -84,7 +84,7 @@ def test_valid_workflow_marks_declarative_workflow_used(self): ], }) - mark_feature_used.assert_called_once_with(FeatureIndex.WORKFLOW) + mark_feature_used.assert_called_once_with(FeatureIndex.DECLARATIVE_WORKFLOW) @_requires_powerfx diff --git a/python/packages/devui/agent_framework_devui/__init__.py b/python/packages/devui/agent_framework_devui/__init__.py index 470134cb09d..a108477d631 100644 --- a/python/packages/devui/agent_framework_devui/__init__.py +++ b/python/packages/devui/agent_framework_devui/__init__.py @@ -8,7 +8,10 @@ from collections.abc import Callable from typing import Any +from agent_framework._telemetry import mark_feature_used + from ._conversations import CheckpointConversationManager +from ._feature_usage import FeatureIndex from ._server import DevServer from .models import AgentFrameworkRequest, OpenAIError, OpenAIResponse, ResponseStreamEvent from .models._discovery_models import DiscoveryResponse, EntityInfo, EnvVarRequirement @@ -196,6 +199,7 @@ def open_browser() -> None: threading.Thread(target=open_browser, daemon=True).start() logger.info(f"Starting Agent Framework DevUI on {host}:{port}") + mark_feature_used(FeatureIndex.DEVUI) uvicorn.run(app, host=host, port=port, log_level="info") diff --git a/python/packages/devui/agent_framework_devui/_feature_usage.py b/python/packages/devui/agent_framework_devui/_feature_usage.py new file mode 100644 index 00000000000..9afe6fa9352 --- /dev/null +++ b/python/packages/devui/agent_framework_devui/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """DevUI-owned feature-usage indexes.""" + + DEVUI = 74 diff --git a/python/packages/durabletask/agent_framework_durabletask/_feature_usage.py b/python/packages/durabletask/agent_framework_durabletask/_feature_usage.py new file mode 100644 index 00000000000..fba5fc35cc0 --- /dev/null +++ b/python/packages/durabletask/agent_framework_durabletask/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Durable Task-owned feature-usage indexes.""" + + DURABLETASK = 77 diff --git a/python/packages/durabletask/agent_framework_durabletask/_shim.py b/python/packages/durabletask/agent_framework_durabletask/_shim.py index e6e9f5d027a..ed8a752a458 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_shim.py +++ b/python/packages/durabletask/agent_framework_durabletask/_shim.py @@ -13,9 +13,11 @@ from typing import Any, Generic, Literal, TypeVar from agent_framework import AgentSession, ServiceSessionId, SupportsAgentRun, normalize_messages +from agent_framework._telemetry import mark_feature_used from agent_framework._types import AgentRunInputs from ._executors import DurableAgentExecutor +from ._feature_usage import FeatureIndex from ._models import DurableAgentSession # TypeVar for the task type returned by executors @@ -127,6 +129,7 @@ def run( # type: ignore[override] options=options, ) + mark_feature_used(FeatureIndex.DURABLETASK) return self._executor.run_durable_agent( agent_name=self.name, run_request=run_request, diff --git a/python/packages/durabletask/agent_framework_durabletask/_worker.py b/python/packages/durabletask/agent_framework_durabletask/_worker.py index 6a3d0193e7a..64bfd543347 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_worker.py +++ b/python/packages/durabletask/agent_framework_durabletask/_worker.py @@ -13,12 +13,14 @@ from typing import Any from agent_framework import SupportsAgentRun, Workflow +from agent_framework._telemetry import mark_feature_used from durabletask.task import ActivityContext, OrchestrationContext from durabletask.worker import TaskHubGrpcWorker from ._async_bridge import run_agent_coroutine from ._callbacks import AgentResponseCallbackProtocol from ._entities import AgentEntity, DurableTaskEntityStateProvider +from ._feature_usage import FeatureIndex from ._workflows.activity import execute_workflow_activity from ._workflows.dt_context import DurableTaskWorkflowContext from ._workflows.naming import ( @@ -157,6 +159,7 @@ def start(self) -> None: The worker will block until stopped. """ logger.info("[DurableAIAgentWorker] Starting worker with %d registered agents", len(self._registered_agents)) + mark_feature_used(FeatureIndex.DURABLETASK) self._worker.start() def stop(self) -> None: diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 21a7e28fa79..7b432c55cc1 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -40,7 +40,7 @@ from ._feature_usage import ( FeatureIndex, - create_feature_usage_user_agent_policy, + create_feature_usage_policy, create_foundry_feature_usage_http_client, ) from ._tools import _sanitize_foundry_response_tool # pyright: ignore[reportPrivateUsage] @@ -179,7 +179,7 @@ class MyClient(FunctionInvocationLayer, RawFoundryAgentChatClient): """ OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" - _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.AGENT + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.FOUNDRY_AGENT def __init__( self, @@ -257,7 +257,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "user_agent_policy": create_feature_usage_user_agent_policy(), + "custom_hook_policy": create_feature_usage_policy(), } if IS_TELEMETRY_ENABLED: project_client_kwargs["user_agent"] = get_user_agent() diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 8f909f004a6..3cb1949147c 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -58,7 +58,7 @@ from ._feature_usage import ( FeatureIndex, - create_feature_usage_user_agent_policy, + create_feature_usage_policy, create_foundry_feature_usage_http_client, ) from ._tools import _sanitize_foundry_response_tool # pyright: ignore[reportPrivateUsage] @@ -156,7 +156,7 @@ class RawFoundryChatClient( OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" SUPPORTS_RICH_FUNCTION_OUTPUT: ClassVar[bool] = False - _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.CHAT_CLIENT + _FEATURE_USAGE_INDEX: ClassVar[int | None] = FeatureIndex.FOUNDRY_CHAT_CLIENT def __init__( self, @@ -226,7 +226,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": project_endpoint, "credential": credential, - "user_agent_policy": create_feature_usage_user_agent_policy(), + "custom_hook_policy": create_feature_usage_policy(), } if IS_TELEMETRY_ENABLED: project_client_kwargs["user_agent"] = get_user_agent() diff --git a/python/packages/foundry/agent_framework_foundry/_embedding_client.py b/python/packages/foundry/agent_framework_foundry/_embedding_client.py index 7e5ed4b4a0c..37ffe4891c2 100644 --- a/python/packages/foundry/agent_framework_foundry/_embedding_client.py +++ b/python/packages/foundry/agent_framework_foundry/_embedding_client.py @@ -23,7 +23,7 @@ from azure.ai.inference.models import ImageEmbeddingInput from azure.core.credentials import AzureKeyCredential -from ._feature_usage import FeatureIndex, create_feature_usage_user_agent_policy +from ._feature_usage import FeatureIndex, create_feature_usage_policy if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover @@ -157,7 +157,7 @@ def __init__( client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "user_agent_policy": create_feature_usage_user_agent_policy(), + "custom_hook_policy": create_feature_usage_policy(), } if IS_TELEMETRY_ENABLED: client_kwargs["user_agent"] = get_user_agent() @@ -210,7 +210,7 @@ async def get_embeddings( """ if not values: return GeneratedEmbeddings([], options=options) - mark_feature_used(FeatureIndex.EMBEDDING) + mark_feature_used(FeatureIndex.FOUNDRY_EMBEDDING) opts: dict[str, Any] = dict(options) if options else {} diff --git a/python/packages/foundry/agent_framework_foundry/_feature_usage.py b/python/packages/foundry/agent_framework_foundry/_feature_usage.py index 260a924abfe..d0f79aaedcb 100644 --- a/python/packages/foundry/agent_framework_foundry/_feature_usage.py +++ b/python/packages/foundry/agent_framework_foundry/_feature_usage.py @@ -4,29 +4,26 @@ from typing import Any from agent_framework._telemetry import ( - IS_TELEMETRY_ENABLED, USER_AGENT_KEY, apply_feature_token, - get_user_agent, remove_feature_token, ) from agent_framework_openai._feature_usage import ( _is_approved_origin, # pyright: ignore[reportPrivateUsage] create_feature_usage_http_client, ) -from azure.core.pipeline.policies import UserAgentPolicy +from azure.core.pipeline.policies import SansIOHTTPPolicy from openai import DefaultAsyncHttpxClient class FeatureIndex(IntEnum): """Foundry-owned feature-usage indexes.""" - CHAT_CLIENT = 48 - AGENT = 49 - MEMORY = 50 - EMBEDDING = 51 - EVALS = 52 - TOOLBOX = 53 + FOUNDRY_CHAT_CLIENT = 48 + FOUNDRY_AGENT = 49 + FOUNDRY_MEMORY = 50 + FOUNDRY_EMBEDDING = 51 + FOUNDRY_EVALS = 52 _FOUNDRY_ORIGIN_SUFFIXES = ( @@ -40,24 +37,22 @@ def create_foundry_feature_usage_http_client() -> DefaultAsyncHttpxClient: return create_feature_usage_http_client(approved_origin_suffixes=_FOUNDRY_ORIGIN_SUFFIXES) -def create_feature_usage_user_agent_policy() -> "FeatureUsageUserAgentPolicy": - """Create the Azure policy with the Agent Framework base User-Agent when enabled.""" - if IS_TELEMETRY_ENABLED: - return FeatureUsageUserAgentPolicy(user_agent=get_user_agent()) - return FeatureUsageUserAgentPolicy() +def create_feature_usage_policy() -> "FeatureUsagePolicy": + """Create the destination-aware policy that stamps each actual request hop.""" + return FeatureUsagePolicy() -class FeatureUsageUserAgentPolicy(UserAgentPolicy): +class FeatureUsagePolicy(SansIOHTTPPolicy[Any, Any]): """Refresh or remove the feature token based on the actual Azure request origin.""" def on_request(self, request: Any) -> None: - """Apply normal Azure User-Agent behavior, then destination-aware feature stamping.""" - super().on_request(request) + """Apply destination-aware feature stamping to the current request hop.""" headers = request.http_request.headers user_agent = headers.get(USER_AGENT_KEY) - base_user_agent = user_agent if isinstance(user_agent, str) else get_user_agent() + if not isinstance(user_agent, str): + return headers[USER_AGENT_KEY] = ( - apply_feature_token(base_user_agent) + apply_feature_token(user_agent) if _is_approved_origin(request.http_request.url, _FOUNDRY_ORIGIN_SUFFIXES) - else remove_feature_token(base_user_agent) + else remove_feature_token(user_agent) ) diff --git a/python/packages/foundry/agent_framework_foundry/_foundry_evals.py b/python/packages/foundry/agent_framework_foundry/_foundry_evals.py index 6fa7cad15ed..16e989b7b02 100644 --- a/python/packages/foundry/agent_framework_foundry/_foundry_evals.py +++ b/python/packages/foundry/agent_framework_foundry/_foundry_evals.py @@ -888,7 +888,7 @@ async def evaluate( Returns: ``EvalResults`` with status, counts, and portal link. """ - mark_feature_used(FeatureIndex.EVALS) + mark_feature_used(FeatureIndex.FOUNDRY_EVALS) # Resolve evaluators with auto-detection resolved = _resolve_default_evaluators(self._evaluators, items=items) # Filter tool evaluators if items don't have tools @@ -1033,7 +1033,7 @@ async def evaluate_traces( ) """ oai_client = _resolve_openai_client(client, project_client) - mark_feature_used(FeatureIndex.EVALS) + mark_feature_used(FeatureIndex.FOUNDRY_EVALS) resolved_evaluators = _resolve_default_evaluators(evaluators) if response_ids: @@ -1126,7 +1126,7 @@ async def evaluate_foundry_target( if "type" not in target: raise ValueError("target dict must include a 'type' key (e.g., 'azure_ai_agent').") oai_client = _resolve_openai_client(client, project_client) - mark_feature_used(FeatureIndex.EVALS) + mark_feature_used(FeatureIndex.FOUNDRY_EVALS) resolved_evaluators = _resolve_default_evaluators(evaluators) eval_obj = await oai_client.evals.create( diff --git a/python/packages/foundry/agent_framework_foundry/_memory_provider.py b/python/packages/foundry/agent_framework_foundry/_memory_provider.py index b8241369be5..ff7f161422e 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -26,7 +26,7 @@ from azure.core.credentials_async import AsyncTokenCredential from openai.types.responses import ResponseInputItemParam -from ._feature_usage import FeatureIndex, create_feature_usage_user_agent_policy +from ._feature_usage import FeatureIndex, create_feature_usage_policy if sys.version_info >= (3, 11): from typing import Self, TypedDict # pragma: no cover @@ -121,7 +121,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "user_agent_policy": create_feature_usage_user_agent_policy(), + "custom_hook_policy": create_feature_usage_policy(), } if IS_TELEMETRY_ENABLED: project_client_kwargs["user_agent"] = get_user_agent() @@ -168,7 +168,7 @@ async def before_run( 2. Searches for contextual memories based on input messages 3. Combines and injects memories into the context """ - mark_feature_used(FeatureIndex.MEMORY) + mark_feature_used(FeatureIndex.FOUNDRY_MEMORY) # On first run, retrieve static memories (user profile memories) if not state.get("initialized"): try: diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 4ab555cb3d6..e25ad14501c 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -45,7 +45,7 @@ _FoundryAgentChatClient, ) from agent_framework_foundry._chat_client import FoundryChatClient -from agent_framework_foundry._feature_usage import FeatureIndex +from agent_framework_foundry._feature_usage import FeatureIndex, FeatureUsagePolicy skip_if_foundry_agent_integration_tests_disabled = pytest.mark.skipif( os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/") @@ -64,7 +64,7 @@ def test_raw_foundry_agent_chat_client_does_not_mark_openai_feature() -> None: assert RawOpenAIChatClient._FEATURE_USAGE_INDEX is OpenAIFeatureIndex.OPENAI - assert RawFoundryAgentChatClient._FEATURE_USAGE_INDEX is FeatureIndex.AGENT + assert RawFoundryAgentChatClient._FEATURE_USAGE_INDEX is FeatureIndex.FOUNDRY_AGENT def _get_foundry_azure_ai_search_model() -> str | None: @@ -126,6 +126,21 @@ def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None: mock_project.get_openai_client.assert_called_once_with(http_client=ANY) +def test_raw_foundry_agent_chat_client_creates_project_client_with_feature_policy() -> None: + mock_project = MagicMock() + mock_project.get_openai_client.return_value = MagicMock() + + with patch("agent_framework_foundry._agent.AIProjectClient", return_value=mock_project) as factory: + RawFoundryAgentChatClient( + project_endpoint="https://test-project.services.ai.azure.com", + credential=MagicMock(), + agent_name="test-agent", + ) + + assert isinstance(factory.call_args.kwargs["custom_hook_policy"], FeatureUsagePolicy) + assert "user_agent_policy" not in factory.call_args.kwargs + + async def test_foundry_agent_basic_call_does_not_request_unsupported_encrypted_reasoning() -> None: """A Foundry agent call must not opt into encrypted reasoning unless the caller requests it.""" mock_response = MagicMock() diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 6d08386137a..17acabc375c 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -20,13 +20,16 @@ from agent_framework_openai._chat_client import RawOpenAIChatClient from azure.ai.projects.models import MCPTool as FoundryMCPTool from azure.core.exceptions import ResourceNotFoundError +from azure.core.pipeline import Pipeline +from azure.core.pipeline.policies import RedirectPolicy, UserAgentPolicy +from azure.core.pipeline.transport import HttpRequest, HttpResponse, HttpTransport from azure.identity import AzureCliCredential from openai import BadRequestError from pydantic import BaseModel from pytest import param from agent_framework_foundry import FoundryChatClient, RawFoundryChatClient -from agent_framework_foundry._feature_usage import FeatureIndex, FeatureUsageUserAgentPolicy +from agent_framework_foundry._feature_usage import FeatureIndex, FeatureUsagePolicy class OutputStruct(BaseModel): @@ -39,19 +42,69 @@ class OutputStruct(BaseModel): def test_foundry_feature_usage_policy_refreshes_user_agent() -> None: with telemetry._feature_mask_lock: telemetry._feature_mask = 0 - mark_feature_used(FeatureIndex.CHAT_CLIENT) + mark_feature_used(FeatureIndex.FOUNDRY_CHAT_CLIENT) request = MagicMock() request.http_request.url = "https://project.services.ai.azure.com/api/projects/test" request.http_request.headers = {"User-Agent": "azsdk-python-ai-projects/1.0 agent-framework-python/1.0"} - FeatureUsageUserAgentPolicy().on_request(request) + FeatureUsagePolicy().on_request(request) assert request.http_request.headers["User-Agent"] == ( "azsdk-python-ai-projects/1.0 agent-framework-python/1.0 (feat=v1.1000000000000)" ) +def test_foundry_feature_usage_policy_removes_token_on_cross_origin_redirect() -> None: + class _Response(HttpResponse): + def body(self) -> bytes: + return b"" + + class _RedirectTransport(HttpTransport[HttpRequest, HttpResponse]): + def __init__(self) -> None: + self.sent_headers: list[dict[str, str]] = [] + + def __enter__(self) -> _RedirectTransport: + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + def open(self) -> None: + pass + + def close(self) -> None: + pass + + def send(self, request: HttpRequest, **kwargs: Any) -> HttpResponse: + self.sent_headers.append(dict(request.headers)) + response = _Response(request, None) + if len(self.sent_headers) == 1: + response.status_code = 302 + response.headers = {"location": "https://example.com/redirected"} + else: + response.status_code = 200 + response.headers = {} + return response + + with telemetry._feature_mask_lock: + telemetry._feature_mask = 0 + mark_feature_used(FeatureIndex.FOUNDRY_CHAT_CLIENT) + transport = _RedirectTransport() + pipeline = cast(Any, Pipeline)( + transport, [UserAgentPolicy(user_agent=get_user_agent()), RedirectPolicy(), FeatureUsagePolicy()] + ) + + pipeline.run(HttpRequest("GET", "https://project.services.ai.azure.com/api/projects/test")) + + assert "(feat=v1." in transport.sent_headers[0]["User-Agent"] + assert "(feat=v1." not in transport.sent_headers[1]["User-Agent"] + + +def test_foundry_feature_index_does_not_own_toolbox() -> None: + assert not hasattr(FeatureIndex, "FOUNDRY_TOOLBOX") + + def test_raw_foundry_chat_client_owns_foundry_feature_bit() -> None: - assert RawFoundryChatClient._FEATURE_USAGE_INDEX is FeatureIndex.CHAT_CLIENT + assert RawFoundryChatClient._FEATURE_USAGE_INDEX is FeatureIndex.FOUNDRY_CHAT_CLIENT @tool(approval_mode="never_require") @@ -238,6 +291,8 @@ def test_init_with_project_endpoint_creates_project_client() -> None: assert factory.call_args.kwargs["credential"] is credential assert factory.call_args.kwargs["allow_preview"] is True assert factory.call_args.kwargs["user_agent"] == get_user_agent() + assert isinstance(factory.call_args.kwargs["custom_hook_policy"], FeatureUsagePolicy) + assert "user_agent_policy" not in factory.call_args.kwargs def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py index 433abee5cfd..ecdd77726d3 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py @@ -211,13 +211,13 @@ def test_settings_from_env(self) -> None: endpoint="https://env.inference.ai.azure.com", credential=ANY, user_agent=get_user_agent(), - user_agent_policy=ANY, + custom_hook_policy=ANY, ) image_client_type.assert_called_once_with( endpoint="https://env.inference.ai.azure.com", credential=ANY, user_agent=get_user_agent(), - user_agent_policy=ANY, + custom_hook_policy=ANY, ) def test_image_model_from_env(self) -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py index 6084161d94a..c5baf53df26 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py +++ b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py @@ -97,7 +97,7 @@ def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMoc credential=mock_credential, allow_preview=True, user_agent=get_user_agent(), - user_agent_policy=ANY, + custom_hook_policy=ANY, ) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_feature_usage.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_feature_usage.py new file mode 100644 index 00000000000..b0e2055c5aa --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_feature_usage.py @@ -0,0 +1,10 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Foundry hosting-owned feature-usage indexes.""" + + FOUNDRY_TOOLBOX = 53 + FOUNDRY_HOSTING = 55 diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index f01c7a94684..4f3ba1f983b 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,12 +1,15 @@ # Copyright (c) Microsoft. All rights reserved. from agent_framework import AgentSession, SupportsAgentRun +from agent_framework._telemetry import mark_feature_used from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.invocations import InvocationAgentServerHost from starlette.requests import Request from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator +from ._feature_usage import FeatureIndex + class InvocationsHostServer(InvocationAgentServerHost): """An invocations server host for an agent.""" @@ -34,6 +37,7 @@ def __init__( self._agent = agent self._sessions: dict[str, AgentSession] = {} self.invoke_handler(self._handle_invoke) + mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) def _partition_key(self) -> str: """Get the partition key for the current request. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 5c9eca58596..ae34e52a1dc 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -26,6 +26,7 @@ SupportsAgentRun, WorkflowAgent, ) +from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentFrameworkException from azure.ai.agentserver.responses import ( ResponseContext, @@ -115,6 +116,8 @@ from mcp import McpError from typing_extensions import Any +from ._feature_usage import FeatureIndex + logger = logging.getLogger(__name__) _AZURE_RESPONSES_MESSAGE_ROLE_TYPE = f"{MessageRole.__module__}:{MessageRole.__qualname__}" @@ -485,6 +488,7 @@ def __init__( self._agent_init_lock = asyncio.Lock() self.shutdown_handler(self._cleanup_agent) self.response_handler(self._handle_response) + mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) async def _ensure_agent_ready(self) -> None: """Lazily enter the agent's async context exactly once. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index 47d96bd4452..9c259a22673 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -19,9 +19,12 @@ SkillsSource, SkillsSourceContext, ) +from agent_framework._telemetry import mark_feature_used from azure.ai.agentserver.core import get_request_context from typing_extensions import override +from ._feature_usage import FeatureIndex + if TYPE_CHECKING: from collections.abc import AsyncGenerator, Generator from datetime import timedelta @@ -218,6 +221,12 @@ def __init__( load_tools=load_tools, ) + @override + async def connect(self, *, reset: bool = False) -> None: + """Connect to the toolbox and mark its first meaningful activation.""" + await super().connect(reset=reset) + mark_feature_used(FeatureIndex.FOUNDRY_TOOLBOX) + @override def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: """Get an authenticated MCP HTTP client. diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 75ba9422543..4ddee21642f 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -9,11 +9,11 @@ from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import cast -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import httpx import pytest -from agent_framework import SkillsProvider, SkillsSourceContext, SupportsAgentRun +from agent_framework import MCPStreamableHTTPTool, SkillsProvider, SkillsSourceContext, SupportsAgentRun from azure.ai.agentserver.core import ( FoundryAgentRequestContext, reset_request_context, @@ -21,6 +21,7 @@ ) from agent_framework_foundry_hosting import FoundryToolbox +from agent_framework_foundry_hosting._feature_usage import FeatureIndex from agent_framework_foundry_hosting._toolbox import ( _FoundryToolboxSkillsSource, _resolve_toolbox_endpoint, @@ -119,6 +120,27 @@ def test_init_derives_name_and_defaults() -> None: assert toolbox.load_prompts_flag is False +def test_toolbox_owns_feature_index_53() -> None: + assert FeatureIndex.FOUNDRY_TOOLBOX == 53 + + +async def test_toolbox_marks_feature_on_successful_connect_not_construction() -> None: + with ( + patch.object(MCPStreamableHTTPTool, "connect", new=AsyncMock()) as connect, + patch("agent_framework_foundry_hosting._toolbox.mark_feature_used") as mark_used, + ): + toolbox = FoundryToolbox( + _FakeCredential(), # type: ignore + url="https://h/toolboxes/sales/mcp?api-version=v1", + ) + mark_used.assert_not_called() + + await toolbox.connect() + + connect.assert_awaited_once_with(reset=False) + mark_used.assert_called_once_with(FeatureIndex.FOUNDRY_TOOLBOX) + + async def test_auth_flow_injects_bearer_token() -> None: cred = _FakeCredential("abc123") auth = _ToolboxAuth(cred, "https://ai.azure.com/.default") # type: ignore diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index b9870544530..cd053711cab 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -29,7 +29,7 @@ validate_tool_mode, ) from agent_framework._settings import SecretString, load_settings -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework._types import _get_data_bytes # type: ignore[reportPrivateUsage] from agent_framework.exceptions import ContentError from agent_framework.observability import ChatTelemetryLayer @@ -38,6 +38,8 @@ from google.genai import types from pydantic import BaseModel +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -539,6 +541,7 @@ def _inner_get_response( async def _stream() -> AsyncIterable[ChatResponseUpdate]: validated = await self._validate_options(options) model, contents, config = self._prepare_request(messages, validated) + mark_feature_used(FeatureIndex.GEMINI) generate_content_stream = cast( Callable[..., Awaitable[AsyncIterable[types.GenerateContentResponse]]], cast(Any, self._genai_client.aio.models).generate_content_stream, @@ -555,6 +558,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: async def _get_response() -> ChatResponse: validated = await self._validate_options(options) model, contents, config = self._prepare_request(messages, validated) + mark_feature_used(FeatureIndex.GEMINI) raw = await self._genai_client.aio.models.generate_content(model=model, contents=contents, config=config) # type: ignore[arg-type] return self._process_generate_response(raw, response_format=validated.get("response_format")) diff --git a/python/packages/gemini/agent_framework_gemini/_feature_usage.py b/python/packages/gemini/agent_framework_gemini/_feature_usage.py new file mode 100644 index 00000000000..fc9da3ff955 --- /dev/null +++ b/python/packages/gemini/agent_framework_gemini/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Gemini-owned feature-usage indexes.""" + + GEMINI = 59 diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index dceb7b72aac..7dbda00ca9d 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -17,6 +17,7 @@ from typing_extensions import NotRequired, TypedDict from agent_framework_gemini import GeminiChatClient, GeminiChatOptions, RawGeminiChatClient, ThinkingConfig +from agent_framework_gemini._feature_usage import FeatureIndex def _has_gemini_integration_credentials() -> bool: @@ -369,8 +370,10 @@ async def test_get_response_returns_text() -> None: client, mock = _make_gemini_client() mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="Hello!")])) - response = await client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])]) + with patch("agent_framework_gemini._chat_client.mark_feature_used") as mark_feature_used: + response = await client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])]) + mark_feature_used.assert_called_once_with(FeatureIndex.GEMINI) assert response.messages[0].text == "Hello!" diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 1a78fb185fe..883810f9fcb 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -29,6 +29,7 @@ normalize_messages, ) from agent_framework._settings import load_settings +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import FunctionTool, ToolTypes from agent_framework._types import ( AgentRunInputs, @@ -38,6 +39,8 @@ from agent_framework.exceptions import AgentException, ContentError from agent_framework.observability import AgentTelemetryLayer +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 11): from typing import Self # pragma: no cover else: @@ -666,6 +669,7 @@ def usage_event_handler(event: SessionEvent) -> None: unsubscribe = copilot_session.on(usage_event_handler) try: + mark_feature_used(FeatureIndex.GITHUB_COPILOT) response_event = await copilot_session.send_and_wait(prompt, attachments=attachments, timeout=timeout) except Exception as ex: raise AgentException(f"GitHub Copilot request failed: {ex}") from ex @@ -852,6 +856,7 @@ def event_handler(event: SessionEvent) -> None: unsubscribe = copilot_session.on(event_handler) try: + mark_feature_used(FeatureIndex.GITHUB_COPILOT) await copilot_session.send(prompt, attachments=attachments) while (item := await queue.get()) is not None: diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_feature_usage.py b/python/packages/github_copilot/agent_framework_github_copilot/_feature_usage.py new file mode 100644 index 00000000000..8facc9c6510 --- /dev/null +++ b/python/packages/github_copilot/agent_framework_github_copilot/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """GitHub Copilot-owned feature-usage indexes.""" + + GITHUB_COPILOT = 64 diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index 5fe74410a93..5e4ff52e05f 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -39,6 +39,7 @@ from copilot.tools import ToolInvocation, ToolResult from agent_framework_github_copilot import GitHubCopilotAgent, GitHubCopilotOptions, RawGitHubCopilotAgent +from agent_framework_github_copilot._feature_usage import FeatureIndex def copilot_options(options: GitHubCopilotOptions) -> GitHubCopilotOptions: @@ -431,8 +432,10 @@ async def test_run_string_message( mock_session.send_and_wait.return_value = assistant_message_event agent = GitHubCopilotAgent(client=mock_client) - response = await agent.run("Hello") + with patch("agent_framework_github_copilot._agent.mark_feature_used") as mark_feature_used: + response = await agent.run("Hello") + mark_feature_used.assert_called_once_with(FeatureIndex.GITHUB_COPILOT) assert isinstance(response, AgentResponse) assert len(response.messages) == 1 assert response.messages[0].role == "assistant" diff --git a/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py b/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py index 583a3e6c62a..3192dda9f0c 100644 --- a/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py +++ b/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py @@ -21,12 +21,15 @@ Workflow, WorkflowRunResult, ) +from agent_framework._telemetry import mark_feature_used from agent_framework_hosting import AgentRunArgs from google.protobuf.json_format import MessageToDict, ParseDict from google.protobuf.struct_pb2 import Value from pydantic import TypeAdapter from pydantic.errors import PydanticSchemaGenerationError +from ._feature_usage import FeatureIndex + logger = logging.getLogger("agent_framework.hosting.a2a") _BINARY_MODE = "application/octet-stream" @@ -189,6 +192,7 @@ def a2a_to_run( ValueError: If the message has no supported content parts or contains a part outside ``input_modes``. """ + mark_feature_used(FeatureIndex.HOSTING_A2A) if input_modes is not None: _validate_part_modes(message.parts, input_modes, "input") @@ -279,6 +283,7 @@ def a2a_from_run( ValueError: If Agent Framework data content contains an invalid data URI or produces a part outside ``output_modes``. """ + mark_feature_used(FeatureIndex.HOSTING_A2A) items: Sequence[Message | AgentResponseUpdate] = result.messages if isinstance(result, AgentResponse) else [result] parts: list[Part] = [] diff --git a/python/packages/hosting-a2a/agent_framework_hosting_a2a/_feature_usage.py b/python/packages/hosting-a2a/agent_framework_hosting_a2a/_feature_usage.py new file mode 100644 index 00000000000..a152324c5e6 --- /dev/null +++ b/python/packages/hosting-a2a/agent_framework_hosting_a2a/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """A2A hosting-owned feature-usage indexes.""" + + HOSTING_A2A = 84 diff --git a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_agent_tool.py b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_agent_tool.py index ac177c0fbda..8905bbda4fa 100644 --- a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_agent_tool.py +++ b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_agent_tool.py @@ -9,10 +9,12 @@ from typing import Any, Generic, TypeVar, cast from agent_framework import AgentResponse, Message, SupportsAgentRun +from agent_framework._telemetry import mark_feature_used from agent_framework_hosting import AgentRunArgs, AgentState from mcp import types from ._conversion import mcp_from_run, mcp_to_run +from ._feature_usage import FeatureIndex AgentT = TypeVar("AgentT", bound=SupportsAgentRun) @@ -86,6 +88,7 @@ def __init__( async def list_tools(self) -> list[types.Tool]: """Return the native MCP tool definition for the target agent.""" + mark_feature_used(FeatureIndex.HOSTING_MCP) target = await self.state.get_target() return [self._tool_for_target(target)] @@ -145,6 +148,7 @@ async def call_tool( Raises: ValueError: If the tool name or configured session id is invalid. """ + mark_feature_used(FeatureIndex.HOSTING_MCP) target = await self.state.get_target() tool = self._tool_for_target(target) if name != tool.name: diff --git a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_feature_usage.py b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_feature_usage.py new file mode 100644 index 00000000000..379c374a10d --- /dev/null +++ b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """MCP hosting-owned feature-usage indexes.""" + + HOSTING_MCP = 85 diff --git a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_workflow_tool.py b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_workflow_tool.py index 6467382e297..b3018e5a31f 100644 --- a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_workflow_tool.py +++ b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_workflow_tool.py @@ -9,11 +9,13 @@ from typing import Any, Generic, TypeVar, cast from agent_framework import AgentResponse, Message, Workflow, WorkflowRunResult +from agent_framework._telemetry import mark_feature_used from agent_framework_hosting import WorkflowState from mcp import types from pydantic import TypeAdapter from ._conversion import mcp_from_run +from ._feature_usage import FeatureIndex WorkflowT = TypeVar("WorkflowT", bound=Workflow) @@ -53,6 +55,7 @@ def __init__( async def list_tools(self) -> list[types.Tool]: """Return the native MCP tool definition for the target workflow.""" + mark_feature_used(FeatureIndex.HOSTING_MCP) workflow = await self.state.get_target() return [self._tool_for_workflow(workflow)] @@ -134,6 +137,7 @@ async def call_tool( Raises: ValueError: If the tool name or workflow input contract is invalid. """ + mark_feature_used(FeatureIndex.HOSTING_MCP) workflow = await self.state.get_target() tool = self._tool_for_workflow(workflow) if name != tool.name: diff --git a/python/packages/hosting-responses/agent_framework_hosting_responses/_feature_usage.py b/python/packages/hosting-responses/agent_framework_hosting_responses/_feature_usage.py new file mode 100644 index 00000000000..df9827b627a --- /dev/null +++ b/python/packages/hosting-responses/agent_framework_hosting_responses/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """OpenAI Responses hosting-owned feature-usage indexes.""" + + HOSTING_RESPONSES = 86 diff --git a/python/packages/hosting-responses/agent_framework_hosting_responses/_parsing.py b/python/packages/hosting-responses/agent_framework_hosting_responses/_parsing.py index 154ce300306..2ced53a629e 100644 --- a/python/packages/hosting-responses/agent_framework_hosting_responses/_parsing.py +++ b/python/packages/hosting-responses/agent_framework_hosting_responses/_parsing.py @@ -21,6 +21,7 @@ from typing import Any, cast from agent_framework import AgentResponse, AgentResponseUpdate, ChatOptions, Content, Message, ResponseStream +from agent_framework._telemetry import mark_feature_used from agent_framework_hosting import AgentRunArgs from openai.types.responses import ( Response as OpenAIResponse, @@ -37,6 +38,8 @@ ) from pydantic import TypeAdapter, ValidationError +from ._feature_usage import FeatureIndex + _RESPONSE_OUTPUT_ITEM_ADAPTER: TypeAdapter[Any] = TypeAdapter(ResponseOutputItem) # OpenAI Responses field name → Agent Framework ChatOptions field name. @@ -179,6 +182,7 @@ def responses_to_run(body: Mapping[str, Any]) -> AgentRunArgs: Raises: ValueError: If the request body has invalid ``input``. """ + mark_feature_used(FeatureIndex.HOSTING_RESPONSES) messages = messages_from_responses_input(body.get("input")) options: dict[str, Any] = {} for key, value in body.items(): @@ -211,6 +215,7 @@ def responses_from_run( Returns: Responses-compatible JSON payload. """ + mark_feature_used(FeatureIndex.HOSTING_RESPONSES) output_items = _result_to_output_items(result, status="completed") response_kwargs: dict[str, Any] = { "id": response_id, @@ -894,6 +899,7 @@ async def responses_from_streaming_run( model: str | None = None updates: list[AgentResponseUpdate] = [] try: + mark_feature_used(FeatureIndex.HOSTING_RESPONSES) async for update in stream: updates.append(update) if model is None: diff --git a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_feature_usage.py b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_feature_usage.py new file mode 100644 index 00000000000..e7a8ae228ba --- /dev/null +++ b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Telegram hosting-owned feature-usage indexes.""" + + HOSTING_TELEGRAM = 87 diff --git a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py index 361d8418676..e738658dad2 100644 --- a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py +++ b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_parsing.py @@ -19,8 +19,11 @@ from typing import Any, cast from agent_framework import ChatOptions, Content, Message +from agent_framework._telemetry import mark_feature_used from agent_framework_hosting import AgentRunArgs +from ._feature_usage import FeatureIndex + # Telegram media fields whose objects carry a `file_id` (and, except photos, # a `mime_type`) directly, mapped to the MIME type Telegram uses when the # object omits `mime_type` (voice notes are always OGG/Opus, for example). @@ -321,6 +324,7 @@ async def telegram_to_run( ValueError: If the update has no actionable message/callback data, or a message has no text, caption, or resolvable media. """ + mark_feature_used(FeatureIndex.HOSTING_TELEGRAM) message = _inner_message(update) if message is not None: contents = await _contents_from_message(message, resolve_file_url) diff --git a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_rendering.py b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_rendering.py index 30d1769096b..3fb963eb121 100644 --- a/python/packages/hosting-telegram/agent_framework_hosting_telegram/_rendering.py +++ b/python/packages/hosting-telegram/agent_framework_hosting_telegram/_rendering.py @@ -15,6 +15,9 @@ from typing import Any, TypedDict from agent_framework import AgentResponse, AgentResponseUpdate, ResponseStream +from agent_framework._telemetry import mark_feature_used + +from ._feature_usage import FeatureIndex # Telegram's documented maximum length, in UTF-16 code units, for message # text (`sendMessage` / `editMessageText`) and photo captions (`sendPhoto`). @@ -81,6 +84,7 @@ def telegram_from_run( Returns: A ``TelegramOperation`` describing the Bot API call to make. """ + mark_feature_used(FeatureIndex.HOSTING_TELEGRAM) text, image_uris = _text_and_image_uris(result) if image_uris: payload: dict[str, Any] = {"chat_id": chat_id, "photo": image_uris[0]} @@ -129,6 +133,7 @@ async def telegram_from_streaming_run( Yields: ``TelegramOperation`` values describing the Bot API calls to make, in order. """ + mark_feature_used(FeatureIndex.HOSTING_TELEGRAM) text = "" last_rendered_text = _truncate(initial_text, TELEGRAM_MAX_TEXT_LENGTH) if initial_text is not None else "" async for update in stream: diff --git a/python/packages/hosting/agent_framework_hosting/_feature_usage.py b/python/packages/hosting/agent_framework_hosting/_feature_usage.py new file mode 100644 index 00000000000..c8ff100238b --- /dev/null +++ b/python/packages/hosting/agent_framework_hosting/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Agent Framework hosting-owned feature-usage indexes.""" + + HOSTING = 83 diff --git a/python/packages/hosting/agent_framework_hosting/_state.py b/python/packages/hosting/agent_framework_hosting/_state.py index 148221e6b44..0f0c82f663b 100644 --- a/python/packages/hosting/agent_framework_hosting/_state.py +++ b/python/packages/hosting/agent_framework_hosting/_state.py @@ -32,6 +32,9 @@ SupportsAgentRun, Workflow, ) +from agent_framework._telemetry import mark_feature_used + +from ._feature_usage import FeatureIndex class SessionStore: @@ -187,6 +190,7 @@ def __init__( self._cached_target = target self._session_store: SessionStore = session_store if session_store is not None else SessionStore() self._session_locks: dict[str, asyncio.Lock] = {} + mark_feature_used(FeatureIndex.HOSTING) async def get_target(self) -> AgentT: """Return the resolved target. @@ -312,6 +316,7 @@ def __init__( self._target_lock = asyncio.Lock() if not callable(target) and not inspect.isawaitable(target): self._cached_target = target + mark_feature_used(FeatureIndex.HOSTING) async def get_target(self) -> WorkflowT: """Return the resolved target. diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_feature_usage.py b/python/packages/hyperlight/agent_framework_hyperlight/_feature_usage.py new file mode 100644 index 00000000000..e97cc56ccc1 --- /dev/null +++ b/python/packages/hyperlight/agent_framework_hyperlight/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Hyperlight-owned feature-usage indexes.""" + + HYPERLIGHT = 81 diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_provider.py b/python/packages/hyperlight/agent_framework_hyperlight/_provider.py index 1232ecc2622..a4fb3a30d6b 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_provider.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_provider.py @@ -7,9 +7,11 @@ from typing import Any from agent_framework import AgentSession, ContextProvider, FunctionTool, SessionContext +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import ApprovalMode from ._execute_code_tool import HyperlightExecuteCodeTool, SandboxRuntime +from ._feature_usage import FeatureIndex from ._types import AllowedDomain, AllowedDomainInput, FileMount, FileMountInput @@ -105,6 +107,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject CodeAct instructions and a run-scoped execute_code tool before each run.""" + mark_feature_used(FeatureIndex.HYPERLIGHT) run_tool = self._execute_code_tool.create_run_tool() state[self.source_id] = run_tool.build_serializable_state() context.extend_instructions(self.source_id, run_tool.build_instructions(tools_visible_to_model=False)) diff --git a/python/packages/lab/common/agent_framework_lab_common/__init__.py b/python/packages/lab/common/agent_framework_lab_common/__init__.py new file mode 100644 index 00000000000..2a50eae8941 --- /dev/null +++ b/python/packages/lab/common/agent_framework_lab_common/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Microsoft. All rights reserved. diff --git a/python/packages/lab/common/agent_framework_lab_common/_feature_usage.py b/python/packages/lab/common/agent_framework_lab_common/_feature_usage.py new file mode 100644 index 00000000000..6f2263768bb --- /dev/null +++ b/python/packages/lab/common/agent_framework_lab_common/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Agent Framework Lab-owned feature-usage indexes.""" + + LAB = 88 diff --git a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py index 5c253b06777..28c9b8a07eb 100644 --- a/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py +++ b/python/packages/lab/gaia/agent_framework_lab_gaia/gaia.py @@ -16,6 +16,8 @@ from pathlib import Path from typing import Any, Protocol, cast +from agent_framework._telemetry import mark_feature_used +from agent_framework_lab_common._feature_usage import FeatureIndex from opentelemetry.trace import NoOpTracer, SpanKind, get_tracer from tqdm import tqdm @@ -201,6 +203,7 @@ def gaia_scorer(model_answer: str | None, ground_truth: str) -> bool: Returns: True if the answer is correct, False otherwise """ + mark_feature_used(FeatureIndex.LAB) def is_float(x: Any) -> bool: try: @@ -393,6 +396,7 @@ def __init__( self.tracer = get_tracer("gaia_benchmark", "1.0.0") else: self.tracer = NoOpTracer() + mark_feature_used(FeatureIndex.LAB) async def _default_evaluator(self, task: Task, prediction: Prediction) -> Evaluation: """Default evaluator using GAIA official scoring.""" @@ -647,6 +651,7 @@ def _save_results(self, results: list[TaskResult], output_path: str) -> None: def viewer_main() -> None: """Main function for the gaia_viewer script.""" + mark_feature_used(FeatureIndex.LAB) import argparse parser = argparse.ArgumentParser(description="View GAIA benchmark results") diff --git a/python/packages/lab/lightning/agent_framework_lab_lightning/__init__.py b/python/packages/lab/lightning/agent_framework_lab_lightning/__init__.py index 688f89825c0..6f35e0b14a9 100644 --- a/python/packages/lab/lightning/agent_framework_lab_lightning/__init__.py +++ b/python/packages/lab/lightning/agent_framework_lab_lightning/__init__.py @@ -6,7 +6,9 @@ import importlib.metadata +from agent_framework._telemetry import mark_feature_used from agent_framework.observability import enable_instrumentation +from agent_framework_lab_common._feature_usage import FeatureIndex from agentlightning.tracer import ( AgentOpsTracer, ) @@ -26,6 +28,7 @@ class AgentFrameworkTracer(AgentOpsTracer): def init(self) -> None: """Initialize the agent-framework-lab-lightning for training.""" + mark_feature_used(FeatureIndex.LAB) enable_instrumentation() super().init() diff --git a/python/packages/lab/pyproject.toml b/python/packages/lab/pyproject.toml index 10d27e4774d..203cf094a60 100644 --- a/python/packages/lab/pyproject.toml +++ b/python/packages/lab/pyproject.toml @@ -87,6 +87,7 @@ packages = [ "agent_framework_lab_gaia", "agent_framework_lab_lightning", "agent_framework_lab_tau2", + "agent_framework_lab_common", "agent_framework.lab.gaia", "agent_framework.lab.lightning", "agent_framework.lab.tau2", @@ -96,6 +97,7 @@ packages = [ "agent_framework_lab_gaia" = "gaia/agent_framework_lab_gaia" "agent_framework_lab_lightning" = "lightning/agent_framework_lab_lightning" "agent_framework_lab_tau2" = "tau2/agent_framework_lab_tau2" +"agent_framework_lab_common" = "common/agent_framework_lab_common" "agent_framework.lab.gaia" = "namespace/agent_framework/lab/gaia" "agent_framework.lab.lightning" = "namespace/agent_framework/lab/lightning" "agent_framework.lab.tau2" = "namespace/agent_framework/lab/tau2" @@ -125,7 +127,12 @@ omit = [ [tool.pyright] extends = "../../pyproject.toml" -include = ["gaia/agent_framework_lab_gaia", "lightning/agent_framework_lab_lightning", "tau2/agent_framework_lab_tau2"] +include = [ + "common/agent_framework_lab_common", + "gaia/agent_framework_lab_gaia", + "lightning/agent_framework_lab_lightning", + "tau2/agent_framework_lab_tau2", +] exclude = ['gaia/tests', 'lightning/tests', 'tau2/tests', 'namespace', '**/samples'] [tool.bandit] diff --git a/python/packages/lab/tau2/agent_framework_lab_tau2/runner.py b/python/packages/lab/tau2/agent_framework_lab_tau2/runner.py index 084c52598de..a4a3657c7cc 100644 --- a/python/packages/lab/tau2/agent_framework_lab_tau2/runner.py +++ b/python/packages/lab/tau2/agent_framework_lab_tau2/runner.py @@ -19,6 +19,8 @@ WorkflowBuilder, WorkflowContext, ) +from agent_framework._telemetry import mark_feature_used +from agent_framework_lab_common._feature_usage import FeatureIndex from loguru import logger from tau2.data_model.simulation import SimulationRun, TerminationReason from tau2.data_model.tasks import Task @@ -338,6 +340,7 @@ async def run( Returns: Complete conversation history as Message list for evaluation """ + mark_feature_used(FeatureIndex.LAB) logger.info(f"Starting workflow agent for task {task.id}: {task.description.purpose}") # type: ignore[unused-ignore] logger.info(f"Assistant chat client: {assistant_chat_client}") logger.info(f"User simulator chat client: {user_simulator_chat_client}") diff --git a/python/packages/mem0/agent_framework_mem0/_context_provider.py b/python/packages/mem0/agent_framework_mem0/_context_provider.py index 15d5945d55d..1ca3138949f 100644 --- a/python/packages/mem0/agent_framework_mem0/_context_provider.py +++ b/python/packages/mem0/agent_framework_mem0/_context_provider.py @@ -17,8 +17,11 @@ from agent_framework import Message from agent_framework._sessions import AgentSession, ContextProvider, SessionContext +from agent_framework._telemetry import mark_feature_used from mem0 import AsyncMemory, AsyncMemoryClient +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 11): from typing import Self # pragma: no cover else: @@ -109,6 +112,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Search Mem0 for relevant memories and add to the session context.""" + mark_feature_used(FeatureIndex.MEM0) self._validate_filters() input_text = "\n".join(msg.text for msg in context.input_messages if msg and msg.text and msg.text.strip()) if not input_text.strip(): @@ -200,6 +204,7 @@ async def after_run( state: dict[str, Any], ) -> None: """Store request/response messages to Mem0 for future retrieval.""" + mark_feature_used(FeatureIndex.MEM0) self._validate_filters() messages_to_store: list[Message] = list(context.input_messages) diff --git a/python/packages/mem0/agent_framework_mem0/_feature_usage.py b/python/packages/mem0/agent_framework_mem0/_feature_usage.py new file mode 100644 index 00000000000..eb422bea73c --- /dev/null +++ b/python/packages/mem0/agent_framework_mem0/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Mem0-owned feature-usage indexes.""" + + MEM0 = 69 diff --git a/python/packages/mem0/tests/test_mem0_context_provider.py b/python/packages/mem0/tests/test_mem0_context_provider.py index 44863b4c119..ab2ddef1fc3 100644 --- a/python/packages/mem0/tests/test_mem0_context_provider.py +++ b/python/packages/mem0/tests/test_mem0_context_provider.py @@ -11,6 +11,7 @@ from agent_framework._sessions import AgentSession, SessionContext from agent_framework_mem0._context_provider import Mem0ContextProvider +from agent_framework_mem0._feature_usage import FeatureIndex @pytest.fixture @@ -600,11 +601,13 @@ async def test_before_run_application_only_fallback(self, mock_mem0_client: Asyn mock_mem0_client.search = AsyncMock(return_value=[{"id": "m1", "memory": "System configuration template"}]) - await provider.before_run( - agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={} - ) + with patch("agent_framework_mem0._context_provider.mark_feature_used") as mark_feature_used: + await provider.before_run( + agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={} + ) # Verify that an application-scoped search task executed successfully + mark_feature_used.assert_called_once_with(FeatureIndex.MEM0) mock_mem0_client.search.assert_awaited_once_with( query="Retrieve systemic fallback memory traces", filters={"app_id": "app_fallback_test"}, diff --git a/python/packages/mistral/agent_framework_mistral/_embedding_client.py b/python/packages/mistral/agent_framework_mistral/_embedding_client.py index 8b76277a0b1..4f6af9e4aa4 100644 --- a/python/packages/mistral/agent_framework_mistral/_embedding_client.py +++ b/python/packages/mistral/agent_framework_mistral/_embedding_client.py @@ -17,8 +17,11 @@ load_settings, ) from agent_framework._settings import SecretString +from agent_framework._telemetry import mark_feature_used from agent_framework.observability import EmbeddingTelemetryLayer +from ._feature_usage import FeatureIndex + def _load_mistral_client_class() -> Any: try: @@ -177,6 +180,7 @@ async def get_embeddings( if "dimensions" in opts: kwargs["output_dimension"] = opts["dimensions"] + mark_feature_used(FeatureIndex.MISTRAL) response = await self.client.embeddings.create_async(**kwargs) embeddings: list[Embedding[list[float]]] = [] diff --git a/python/packages/mistral/agent_framework_mistral/_feature_usage.py b/python/packages/mistral/agent_framework_mistral/_feature_usage.py new file mode 100644 index 00000000000..480733ad74f --- /dev/null +++ b/python/packages/mistral/agent_framework_mistral/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Mistral-owned feature-usage indexes.""" + + MISTRAL = 60 diff --git a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py index cfc9e4e72e7..09f47973a5f 100644 --- a/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py +++ b/python/packages/mistral/tests/mistral/test_mistral_embedding_client.py @@ -10,6 +10,7 @@ from agent_framework_mistral import MistralEmbeddingClient, MistralEmbeddingOptions from agent_framework_mistral._embedding_client import _load_mistral_client_class # pyright: ignore[reportPrivateUsage] +from agent_framework_mistral._feature_usage import FeatureIndex # region: Unit Tests @@ -137,7 +138,10 @@ async def test_mistral_embedding_get_embeddings() -> None: mock_response.model = "mistral-embed" mock_response.usage = MagicMock(prompt_tokens=10, total_tokens=10) - with patch("agent_framework_mistral._embedding_client.Mistral") as mock_cls: + with ( + patch("agent_framework_mistral._embedding_client.Mistral") as mock_cls, + patch("agent_framework_mistral._embedding_client.mark_feature_used") as mark_feature_used, + ): mock_client = MagicMock() mock_client.embeddings = MagicMock() mock_client.embeddings.create_async = AsyncMock(return_value=mock_response) @@ -146,6 +150,7 @@ async def test_mistral_embedding_get_embeddings() -> None: client = MistralEmbeddingClient(model="mistral-embed", api_key="test-key") result = await client.get_embeddings(["hello", "world"]) + mark_feature_used.assert_called_once_with(FeatureIndex.MISTRAL) assert isinstance(result, GeneratedEmbeddings) assert len(result) == 2 assert result[0].vector == [0.1, 0.2, 0.3] diff --git a/python/packages/monty/agent_framework_monty/_feature_usage.py b/python/packages/monty/agent_framework_monty/_feature_usage.py new file mode 100644 index 00000000000..03f71b545df --- /dev/null +++ b/python/packages/monty/agent_framework_monty/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Monty-owned feature-usage indexes.""" + + MONTY = 80 diff --git a/python/packages/monty/agent_framework_monty/_provider.py b/python/packages/monty/agent_framework_monty/_provider.py index abec2a33fab..6e45e958a19 100644 --- a/python/packages/monty/agent_framework_monty/_provider.py +++ b/python/packages/monty/agent_framework_monty/_provider.py @@ -9,9 +9,11 @@ from typing import Any from agent_framework import AgentSession, ContextProvider, FunctionTool, SessionContext +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import ApprovalMode from ._execute_code_tool import MontyExecuteCodeTool +from ._feature_usage import FeatureIndex from ._types import FileMount, FileMountInput @@ -89,6 +91,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Inject CodeAct instructions and a run-scoped execute_code tool before each run.""" + mark_feature_used(FeatureIndex.MONTY) run_tool = self._execute_code_tool.create_run_tool() state[self.source_id] = run_tool.build_serializable_state() context.extend_instructions(self.source_id, run_tool.build_instructions(tools_visible_to_model=False)) diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index 94e612f75db..f8c7a2232ec 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -33,6 +33,7 @@ UsageDetails, ) from agent_framework._settings import load_settings +from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import ( ChatClientException, ChatClientInvalidRequestException, @@ -45,6 +46,8 @@ from ollama._types import Message as OllamaMessage from pydantic import BaseModel +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -358,6 +361,7 @@ def _inner_get_response( async def _stream() -> AsyncIterable[ChatResponseUpdate]: validated_options = await self._validate_options(options) options_dict = self._prepare_options(messages, validated_options) + mark_feature_used(FeatureIndex.OLLAMA) try: response_object: AsyncIterable[OllamaChatResponse] = await self.client.chat( # type: ignore[misc] stream=True, @@ -376,6 +380,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: async def _get_response() -> ChatResponse: validated_options = await self._validate_options(options) options_dict = self._prepare_options(messages, validated_options) + mark_feature_used(FeatureIndex.OLLAMA) try: response: OllamaChatResponse = await self.client.chat( # type: ignore[misc] stream=False, diff --git a/python/packages/ollama/agent_framework_ollama/_embedding_client.py b/python/packages/ollama/agent_framework_ollama/_embedding_client.py index cb8a004efee..03f7078c5c6 100644 --- a/python/packages/ollama/agent_framework_ollama/_embedding_client.py +++ b/python/packages/ollama/agent_framework_ollama/_embedding_client.py @@ -15,9 +15,12 @@ UsageDetails, load_settings, ) +from agent_framework._telemetry import mark_feature_used from agent_framework.observability import EmbeddingTelemetryLayer from ollama import AsyncClient +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 13): from typing import TypeVar # pragma: no cover else: @@ -150,6 +153,7 @@ async def get_embeddings( if dimensions := opts.get("dimensions"): kwargs["dimensions"] = dimensions + mark_feature_used(FeatureIndex.OLLAMA) response = await self.client.embed(**kwargs) embeddings = [ diff --git a/python/packages/ollama/agent_framework_ollama/_feature_usage.py b/python/packages/ollama/agent_framework_ollama/_feature_usage.py new file mode 100644 index 00000000000..8b83ca44d26 --- /dev/null +++ b/python/packages/ollama/agent_framework_ollama/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Ollama-owned feature-usage indexes.""" + + OLLAMA = 61 diff --git a/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py b/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py index 63dffeb7c19..9cd12bbc3f8 100644 --- a/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py +++ b/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py @@ -7,6 +7,7 @@ from agent_framework import Embedding, GeneratedEmbeddings from agent_framework_ollama import OllamaEmbeddingClient, OllamaEmbeddingOptions +from agent_framework_ollama._feature_usage import FeatureIndex # region: Unit Tests @@ -49,7 +50,10 @@ async def test_ollama_embedding_get_embeddings() -> None: "prompt_eval_count": 10, } - with patch("agent_framework_ollama._embedding_client.AsyncClient") as mock_client_cls: + with ( + patch("agent_framework_ollama._embedding_client.AsyncClient") as mock_client_cls, + patch("agent_framework_ollama._embedding_client.mark_feature_used") as mark_feature_used, + ): mock_client = MagicMock() mock_client.embed = AsyncMock(return_value=mock_response) mock_client_cls.return_value = mock_client @@ -57,6 +61,7 @@ async def test_ollama_embedding_get_embeddings() -> None: client = OllamaEmbeddingClient(model="nomic-embed-text") result = await client.get_embeddings(["hello", "world"]) + mark_feature_used.assert_called_once_with(FeatureIndex.OLLAMA) assert isinstance(result, GeneratedEmbeddings) assert len(result) == 2 assert result[0].vector == [0.1, 0.2, 0.3] diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py b/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py index 9c5c8b2ab05..5d9e20bdaed 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_concurrent.py @@ -404,7 +404,7 @@ def build(self) -> Workflow: workflow = ConcurrentBuilder(participants=[agent1, agent2]).build() """ - mark_feature_used(FeatureIndex.CONCURRENT) + mark_feature_used(FeatureIndex.ORCHESTRATION_CONCURRENT) # Internal nodes dispatcher = _DispatchToAllParticipants(id="dispatcher") aggregator = self._aggregator if self._aggregator is not None else _AggregateAgentConversations(id="aggregator") diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py b/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py index 86d804a061b..66e9c620173 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_feature_usage.py @@ -6,8 +6,8 @@ class FeatureIndex(IntEnum): """Orchestration-owned feature-usage indexes.""" - SEQUENTIAL = 32 - CONCURRENT = 33 - GROUP_CHAT = 34 - MAGENTIC = 35 - HANDOFF = 36 + ORCHESTRATION_SEQUENTIAL = 32 + ORCHESTRATION_CONCURRENT = 33 + ORCHESTRATION_GROUP_CHAT = 34 + ORCHESTRATION_MAGENTIC = 35 + ORCHESTRATION_HANDOFF = 36 diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py index d3cac2158fd..bbb61edae28 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_group_chat.py @@ -1012,7 +1012,7 @@ def build(self) -> Workflow: Returns: Validated Workflow instance ready for execution """ - mark_feature_used(FeatureIndex.GROUP_CHAT) + mark_feature_used(FeatureIndex.ORCHESTRATION_GROUP_CHAT) # Resolve orchestrator and participants to executors participants: list[Executor] = self._resolve_participants() orchestrator: Executor = self._resolve_orchestrator(participants) diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py b/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py index 04d40432c72..4c0de8001be 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py @@ -943,7 +943,7 @@ def build(self) -> Workflow: ValueError: If participants or coordinator were not configured, or if required configuration is invalid. """ - mark_feature_used(FeatureIndex.HANDOFF) + mark_feature_used(FeatureIndex.ORCHESTRATION_HANDOFF) # Resolve agents (either from instances or factories) # The returned map keys are either executor IDs or factory names, which is need to resolve handoff configs resolved_agents = self._resolve_agents() diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py index 902b26d652c..c5406001bd1 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py @@ -1776,7 +1776,7 @@ def _resolve_participants(self) -> list[Executor]: def build(self) -> Workflow: """Build a Magentic workflow with the orchestrator and all agent executors.""" - mark_feature_used(FeatureIndex.MAGENTIC) + mark_feature_used(FeatureIndex.ORCHESTRATION_MAGENTIC) logger.info(f"Building Magentic workflow with {len(self._participants)} participants") participants: list[Executor] = self._resolve_participants() diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py b/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py index 4da4a3a7c20..fbb98179501 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_sequential.py @@ -243,7 +243,7 @@ def build(self) -> Workflow: terminator's own `yield_output` is Workflow Output (`AgentResponse`, or per-chunk `AgentResponseUpdate` when streaming). """ - mark_feature_used(FeatureIndex.SEQUENTIAL) + mark_feature_used(FeatureIndex.ORCHESTRATION_SEQUENTIAL) input_conv = _InputToConversation(id="input-conversation") # Resolve participants and participant factories to executors diff --git a/python/packages/orchestrations/tests/test_magentic.py b/python/packages/orchestrations/tests/test_magentic.py index ec586a1ca3a..a759c354400 100644 --- a/python/packages/orchestrations/tests/test_magentic.py +++ b/python/packages/orchestrations/tests/test_magentic.py @@ -198,7 +198,7 @@ def test_magentic_builder_marks_feature_with_custom_manager() -> None: token = get_feature_token() assert token is not None - assert int(token.split(".", 1)[1], 16) & (1 << FeatureIndex.MAGENTIC) + assert int(token.split(".", 1)[1], 16) & (1 << FeatureIndex.ORCHESTRATION_MAGENTIC) async def test_magentic_builder_returns_workflow_and_runs() -> None: diff --git a/python/packages/orchestrations/tests/test_sequential.py b/python/packages/orchestrations/tests/test_sequential.py index 758762ea89f..92dc7d409da 100644 --- a/python/packages/orchestrations/tests/test_sequential.py +++ b/python/packages/orchestrations/tests/test_sequential.py @@ -105,7 +105,7 @@ def test_sequential_builder_does_not_mark_custom_workflow() -> None: token = get_feature_token() assert token is not None mask = int(token.split(".", 1)[1], 16) - assert mask & (1 << FeatureIndex.SEQUENTIAL) + assert mask & (1 << FeatureIndex.ORCHESTRATION_SEQUENTIAL) assert not mask & (1 << CoreFeatureIndex.CORE_WORKFLOW) diff --git a/python/packages/purview/agent_framework_purview/_client.py b/python/packages/purview/agent_framework_purview/_client.py index 9208d02f557..eb6ed39381c 100644 --- a/python/packages/purview/agent_framework_purview/_client.py +++ b/python/packages/purview/agent_framework_purview/_client.py @@ -11,7 +11,7 @@ from uuid import uuid4 import httpx -from agent_framework._telemetry import get_user_agent +from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework.observability import get_tracer from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -24,6 +24,7 @@ PurviewRequestError, PurviewServiceError, ) +from ._feature_usage import FeatureIndex from ._models import ( ContentActivitiesRequest, ContentActivitiesResponse, @@ -96,10 +97,12 @@ def _extract_token_info(token: str) -> dict[str, Any]: } async def get_user_info_from_token(self, *, tenant_id: str | None = None) -> dict[str, Any]: + mark_feature_used(FeatureIndex.PURVIEW) token = await self._get_token(tenant_id=tenant_id) return self._extract_token_info(token) async def process_content(self, request: ProcessContentRequest) -> ProcessContentResponse: + mark_feature_used(FeatureIndex.PURVIEW) with get_tracer().start_as_current_span("purview.process_content"): token = await self._get_token(tenant_id=request.tenant_id) url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/processContent" @@ -122,6 +125,7 @@ async def process_content(self, request: ProcessContentRequest) -> ProcessConten return response async def get_protection_scopes(self, request: ProtectionScopesRequest) -> ProtectionScopesResponse: + mark_feature_used(FeatureIndex.PURVIEW) with get_tracer().start_as_current_span("purview.get_protection_scopes"): token = await self._get_token() url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/protectionScopes/compute" @@ -140,6 +144,7 @@ async def get_protection_scopes(self, request: ProtectionScopesRequest) -> Prote return response async def send_content_activities(self, request: ContentActivitiesRequest) -> ContentActivitiesResponse: + mark_feature_used(FeatureIndex.PURVIEW) with get_tracer().start_as_current_span("purview.send_content_activities"): token = await self._get_token() url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/activities/contentActivities" diff --git a/python/packages/purview/agent_framework_purview/_feature_usage.py b/python/packages/purview/agent_framework_purview/_feature_usage.py new file mode 100644 index 00000000000..0d9d71cfdf9 --- /dev/null +++ b/python/packages/purview/agent_framework_purview/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Purview-owned feature-usage indexes.""" + + PURVIEW = 70 diff --git a/python/packages/purview/tests/purview/test_purview_client.py b/python/packages/purview/tests/purview/test_purview_client.py index 133194f1824..b4c29b15312 100644 --- a/python/packages/purview/tests/purview/test_purview_client.py +++ b/python/packages/purview/tests/purview/test_purview_client.py @@ -19,6 +19,7 @@ PurviewRequestError, PurviewServiceError, ) +from agent_framework_purview._feature_usage import FeatureIndex from agent_framework_purview._models import ( ContentActivitiesRequest, ContentActivitiesResponse, @@ -167,9 +168,13 @@ async def test_process_content_success( mock_response.headers = {} mock_response.json.return_value = {"id": "response-123", "protectionScopeState": "notModified"} - with patch.object(client._client, "post", return_value=mock_response): + with ( + patch.object(client._client, "post", return_value=mock_response), + patch("agent_framework_purview._client.mark_feature_used") as mark_feature_used, + ): response = await client.process_content(request) + mark_feature_used.assert_called_once_with(FeatureIndex.PURVIEW) assert response.id == "response-123" assert response.protection_scope_state == "notModified" diff --git a/python/packages/redis/agent_framework_redis/_context_provider.py b/python/packages/redis/agent_framework_redis/_context_provider.py index b3ac7508179..44c9c675abd 100644 --- a/python/packages/redis/agent_framework_redis/_context_provider.py +++ b/python/packages/redis/agent_framework_redis/_context_provider.py @@ -15,6 +15,7 @@ import numpy as np from agent_framework import Message from agent_framework._sessions import AgentSession, ContextProvider, SessionContext +from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import ( AgentException, IntegrationInvalidRequestException, @@ -25,6 +26,8 @@ from redisvl.utils.token_escaper import TokenEscaper from redisvl.utils.vectorize import BaseVectorizer +from ._feature_usage import FeatureIndex + if sys.version_info >= (3, 11): from typing import Self # pragma: no cover else: @@ -122,6 +125,7 @@ async def before_run( state: dict[str, Any], ) -> None: """Retrieve scoped context from Redis and add to the session context.""" + mark_feature_used(FeatureIndex.REDIS) self._validate_filters() input_text = "\n".join(msg.text for msg in context.input_messages if msg and msg.text and msg.text.strip()) if not input_text.strip(): @@ -147,6 +151,7 @@ async def after_run( state: dict[str, Any], ) -> None: """Store request/response messages to Redis for future retrieval.""" + mark_feature_used(FeatureIndex.REDIS) self._validate_filters() messages_to_store: list[Message] = list(context.input_messages) diff --git a/python/packages/redis/agent_framework_redis/_feature_usage.py b/python/packages/redis/agent_framework_redis/_feature_usage.py new file mode 100644 index 00000000000..54b9e42799a --- /dev/null +++ b/python/packages/redis/agent_framework_redis/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Redis-owned feature-usage indexes.""" + + REDIS = 68 diff --git a/python/packages/redis/agent_framework_redis/_history_provider.py b/python/packages/redis/agent_framework_redis/_history_provider.py index ef22511faa8..a7703db8a2c 100644 --- a/python/packages/redis/agent_framework_redis/_history_provider.py +++ b/python/packages/redis/agent_framework_redis/_history_provider.py @@ -14,8 +14,11 @@ import redis.asyncio as redis from agent_framework import Message from agent_framework._sessions import HistoryProvider +from agent_framework._telemetry import mark_feature_used from redis.credentials import CredentialProvider +from ._feature_usage import FeatureIndex + class RedisHistoryProvider(HistoryProvider): """Redis-backed history provider using the new HistoryProvider hooks pattern. @@ -124,6 +127,7 @@ async def get_messages( Returns: List of stored Message objects in chronological order. """ + mark_feature_used(FeatureIndex.REDIS) key = self._redis_key(session_id) redis_messages: list[str] = await self._redis_client.lrange(key, 0, -1) # type: ignore[misc] messages: list[Message] = [] @@ -148,6 +152,7 @@ async def save_messages( state: Optional session state. Unused for Redis-backed history. **kwargs: Additional arguments (unused). """ + mark_feature_used(FeatureIndex.REDIS) if not messages: return diff --git a/python/packages/redis/tests/test_providers.py b/python/packages/redis/tests/test_providers.py index 74ed328b404..55aee296626 100644 --- a/python/packages/redis/tests/test_providers.py +++ b/python/packages/redis/tests/test_providers.py @@ -13,6 +13,7 @@ from agent_framework._sessions import AgentSession, SessionContext from agent_framework_redis._context_provider import RedisContextProvider +from agent_framework_redis._feature_usage import FeatureIndex from agent_framework_redis._history_provider import RedisHistoryProvider # --------------------------------------------------------------------------- @@ -20,6 +21,15 @@ # --------------------------------------------------------------------------- +async def test_empty_history_save_marks_redis_used() -> None: + provider = object.__new__(RedisHistoryProvider) + + with patch("agent_framework_redis._history_provider.mark_feature_used") as mark_feature_used: + await provider.save_messages(None, []) + + mark_feature_used.assert_called_once_with(FeatureIndex.REDIS) + + @pytest.fixture def mock_index() -> AsyncMock: idx = AsyncMock() diff --git a/python/packages/tools/agent_framework_tools/_feature_usage.py b/python/packages/tools/agent_framework_tools/_feature_usage.py new file mode 100644 index 00000000000..4805f13ec31 --- /dev/null +++ b/python/packages/tools/agent_framework_tools/_feature_usage.py @@ -0,0 +1,9 @@ +# Copyright (c) Microsoft. All rights reserved. + +from enum import IntEnum + + +class FeatureIndex(IntEnum): + """Agent Framework tools-owned feature-usage indexes.""" + + TOOLS_SHELL = 79 diff --git a/python/packages/tools/agent_framework_tools/shell/_docker.py b/python/packages/tools/agent_framework_tools/shell/_docker.py index 189509dc159..c83aa1318f3 100644 --- a/python/packages/tools/agent_framework_tools/shell/_docker.py +++ b/python/packages/tools/agent_framework_tools/shell/_docker.py @@ -53,8 +53,10 @@ from typing import Literal from agent_framework import FunctionTool, tool +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import SHELL_TOOL_KIND_VALUE +from .._feature_usage import FeatureIndex from ._policy import ShellPolicy, ShellRequest from ._session import ShellSession from ._truncate import truncate_head_tail as _truncate_bytes @@ -445,6 +447,7 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul caller does not need to wrap the call in :func:`asyncio.wait_for`. """ + mark_feature_used(FeatureIndex.TOOLS_SHELL) request = ShellRequest(command=command, workdir=self._workdir) decision = self._policy.evaluate(request) if decision.decision == "deny": diff --git a/python/packages/tools/agent_framework_tools/shell/_tool.py b/python/packages/tools/agent_framework_tools/shell/_tool.py index 920639c0737..55acd1af7ce 100644 --- a/python/packages/tools/agent_framework_tools/shell/_tool.py +++ b/python/packages/tools/agent_framework_tools/shell/_tool.py @@ -11,8 +11,10 @@ from typing import Literal from agent_framework import FunctionTool, tool +from agent_framework._telemetry import mark_feature_used from agent_framework._tools import SHELL_TOOL_KIND_VALUE +from .._feature_usage import FeatureIndex from ._executor import run_stateless from ._policy import ShellPolicy, ShellRequest from ._resolve import is_powershell, resolve_shell @@ -244,6 +246,7 @@ async def run(self, command: str, *, timeout: float | None = None) -> ShellResul so callers do not need to wrap this call in :func:`asyncio.wait_for`. """ + mark_feature_used(FeatureIndex.TOOLS_SHELL) request = ShellRequest(command=command, workdir=self._workdir) decision = self._policy.evaluate(request) if decision.decision == "deny": From 2dcc30113d4a617b39253099342cf63b5689bd8a Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 30 Jul 2026 07:50:45 +0200 Subject: [PATCH 4/8] Python: report core version in User-Agent Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- python/packages/core/agent_framework/__init__.py | 2 +- python/packages/core/tests/core/test_telemetry.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 45ac531dd6c..f1bffab9b1e 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -18,7 +18,7 @@ from typing import Any, Final try: - _version = importlib.metadata.version(__name__) + _version = importlib.metadata.version("agent-framework-core") except importlib.metadata.PackageNotFoundError: _version = "0.0.0" # Fallback for development mode __version__: Final[str] = _version diff --git a/python/packages/core/tests/core/test_telemetry.py b/python/packages/core/tests/core/test_telemetry.py index 3add7e5c993..80bb029769e 100644 --- a/python/packages/core/tests/core/test_telemetry.py +++ b/python/packages/core/tests/core/test_telemetry.py @@ -2,6 +2,7 @@ import ast import concurrent.futures +import importlib.metadata import os import re from pathlib import Path @@ -48,6 +49,12 @@ def test_agent_framework_user_agent_format(): assert AGENT_FRAMEWORK_USER_AGENT.startswith("agent-framework-python/") +def test_agent_framework_user_agent_uses_core_distribution_version() -> None: + core_version = importlib.metadata.version("agent-framework-core") + + assert f"agent-framework-python/{core_version}" == AGENT_FRAMEWORK_USER_AGENT + + def _reset_feature_mask() -> None: with _telemetry_mod._feature_mask_lock: _telemetry_mod._feature_mask = 0 From d58ce6acd897b7e8f38e44114021d35b53fd2b38 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 30 Jul 2026 07:54:59 +0200 Subject: [PATCH 5/8] Python: configure Lab telemetry import path Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- python/packages/lab/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/python/packages/lab/pyproject.toml b/python/packages/lab/pyproject.toml index 203cf094a60..e600798bc84 100644 --- a/python/packages/lab/pyproject.toml +++ b/python/packages/lab/pyproject.toml @@ -133,6 +133,7 @@ include = [ "lightning/agent_framework_lab_lightning", "tau2/agent_framework_lab_tau2", ] +extraPaths = ["common"] exclude = ['gaia/tests', 'lightning/tests', 'tau2/tests', 'namespace', '**/samples'] [tool.bandit] From ab4326f962280db2647e16cab46b5b7c823e8ee6 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 30 Jul 2026 08:12:54 +0200 Subject: [PATCH 6/8] Python: preserve telemetry transport behavior Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- .../foundry/agent_framework_foundry/_agent.py | 2 +- .../agent_framework_foundry/_chat_client.py | 2 +- .../agent_framework_foundry/_embedding_client.py | 11 ++++++++--- .../agent_framework_foundry/_memory_provider.py | 2 +- .../foundry/tests/foundry/test_foundry_agent.py | 6 ++++-- .../tests/foundry/test_foundry_chat_client.py | 6 ++++-- .../foundry/test_foundry_embedding_client.py | 4 ++-- .../tests/foundry/test_foundry_memory_provider.py | 2 +- .../agent_framework_openai/_feature_usage.py | 14 +++++++++++++- .../openai/test_openai_chat_completion_client.py | 3 +++ .../openai/tests/openai/test_openai_shared.py | 15 ++++++++++++++- 11 files changed, 52 insertions(+), 15 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 7b432c55cc1..61b892da06a 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -257,7 +257,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "custom_hook_policy": create_feature_usage_policy(), + "per_retry_policies": [create_feature_usage_policy()], } if IS_TELEMETRY_ENABLED: project_client_kwargs["user_agent"] = get_user_agent() diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 3cb1949147c..6082aecb54b 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -226,7 +226,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": project_endpoint, "credential": credential, - "custom_hook_policy": create_feature_usage_policy(), + "per_retry_policies": [create_feature_usage_policy()], } if IS_TELEMETRY_ENABLED: project_client_kwargs["user_agent"] = get_user_agent() diff --git a/python/packages/foundry/agent_framework_foundry/_embedding_client.py b/python/packages/foundry/agent_framework_foundry/_embedding_client.py index 37ffe4891c2..779738e9150 100644 --- a/python/packages/foundry/agent_framework_foundry/_embedding_client.py +++ b/python/packages/foundry/agent_framework_foundry/_embedding_client.py @@ -157,12 +157,17 @@ def __init__( client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "custom_hook_policy": create_feature_usage_policy(), } if IS_TELEMETRY_ENABLED: client_kwargs["user_agent"] = get_user_agent() - self._text_client = text_client or EmbeddingsClient(**client_kwargs) - self._image_client = image_client or ImageEmbeddingsClient(**client_kwargs) + self._text_client = text_client or EmbeddingsClient( + **client_kwargs, + per_retry_policies=[create_feature_usage_policy()], + ) + self._image_client = image_client or ImageEmbeddingsClient( + **client_kwargs, + per_retry_policies=[create_feature_usage_policy()], + ) self._endpoint = resolved_endpoint super().__init__(additional_properties=additional_properties) diff --git a/python/packages/foundry/agent_framework_foundry/_memory_provider.py b/python/packages/foundry/agent_framework_foundry/_memory_provider.py index ff7f161422e..6f65d8727e4 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -121,7 +121,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "custom_hook_policy": create_feature_usage_policy(), + "per_retry_policies": [create_feature_usage_policy()], } if IS_TELEMETRY_ENABLED: project_client_kwargs["user_agent"] = get_user_agent() diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index e25ad14501c..c8473fd3689 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -137,8 +137,10 @@ def test_raw_foundry_agent_chat_client_creates_project_client_with_feature_polic agent_name="test-agent", ) - assert isinstance(factory.call_args.kwargs["custom_hook_policy"], FeatureUsagePolicy) - assert "user_agent_policy" not in factory.call_args.kwargs + policies = factory.call_args.kwargs["per_retry_policies"] + assert len(policies) == 1 + assert isinstance(policies[0], FeatureUsagePolicy) + assert "custom_hook_policy" not in factory.call_args.kwargs async def test_foundry_agent_basic_call_does_not_request_unsupported_encrypted_reasoning() -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 17acabc375c..128fbb0e04d 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -291,8 +291,10 @@ def test_init_with_project_endpoint_creates_project_client() -> None: assert factory.call_args.kwargs["credential"] is credential assert factory.call_args.kwargs["allow_preview"] is True assert factory.call_args.kwargs["user_agent"] == get_user_agent() - assert isinstance(factory.call_args.kwargs["custom_hook_policy"], FeatureUsagePolicy) - assert "user_agent_policy" not in factory.call_args.kwargs + policies = factory.call_args.kwargs["per_retry_policies"] + assert len(policies) == 1 + assert isinstance(policies[0], FeatureUsagePolicy) + assert "custom_hook_policy" not in factory.call_args.kwargs def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py index ecdd77726d3..827f2ea4343 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_embedding_client.py @@ -211,13 +211,13 @@ def test_settings_from_env(self) -> None: endpoint="https://env.inference.ai.azure.com", credential=ANY, user_agent=get_user_agent(), - custom_hook_policy=ANY, + per_retry_policies=[ANY], ) image_client_type.assert_called_once_with( endpoint="https://env.inference.ai.azure.com", credential=ANY, user_agent=get_user_agent(), - custom_hook_policy=ANY, + per_retry_policies=[ANY], ) def test_image_model_from_env(self) -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py index c5baf53df26..5d7e16d0f9b 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py +++ b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py @@ -97,7 +97,7 @@ def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMoc credential=mock_credential, allow_preview=True, user_agent=get_user_agent(), - custom_hook_policy=ANY, + per_retry_policies=[ANY], ) diff --git a/python/packages/openai/agent_framework_openai/_feature_usage.py b/python/packages/openai/agent_framework_openai/_feature_usage.py index 59208b0b96e..7b92d1aa3dc 100644 --- a/python/packages/openai/agent_framework_openai/_feature_usage.py +++ b/python/packages/openai/agent_framework_openai/_feature_usage.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio +import contextlib from enum import IntEnum import httpx @@ -20,6 +22,16 @@ class FeatureIndex(IntEnum): ) +class _FeatureUsageAsyncHttpxClient(DefaultAsyncHttpxClient): + """OpenAI-default HTTP client that preserves the SDK's GC cleanup behavior.""" + + def __del__(self) -> None: + if self.is_closed: + return + with contextlib.suppress(Exception): + asyncio.get_running_loop().create_task(self.aclose()) + + def _is_approved_origin(url: httpx.URL | str, suffixes: tuple[str, ...]) -> bool: if isinstance(url, str): url = httpx.URL(url) @@ -41,4 +53,4 @@ async def stamp_feature_usage(request: httpx.Request) -> None: # ruff:ignore[un else remove_feature_token(user_agent) ) - return DefaultAsyncHttpxClient(event_hooks={"request": [stamp_feature_usage]}) + return _FeatureUsageAsyncHttpxClient(event_hooks={"request": [stamp_feature_usage]}) diff --git a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py index f769b46716e..d933c280f0b 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_completion_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_completion_client.py @@ -1921,6 +1921,9 @@ async def chunks() -> Any: options={}, ) assert isinstance(stream, ResponseStream) + token = telemetry.get_feature_token() + assert token is None or not int(token.split(".", 1)[1], 16) & (1 << FeatureIndex.OPENAI) + mark_feature_used(CoreFeatureIndex.CORE_AGENT) async for _ in stream: pass diff --git a/python/packages/openai/tests/openai/test_openai_shared.py b/python/packages/openai/tests/openai/test_openai_shared.py index e485c55a62b..f844f3f90d0 100644 --- a/python/packages/openai/tests/openai/test_openai_shared.py +++ b/python/packages/openai/tests/openai/test_openai_shared.py @@ -2,9 +2,10 @@ from __future__ import annotations +import asyncio from types import TracebackType from typing import Any, cast -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import agent_framework._telemetry as telemetry import httpx @@ -102,6 +103,18 @@ async def test_feature_usage_hook_stamps_approved_origin_and_strips_custom_origi assert custom.headers["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT +async def test_feature_usage_http_client_preserves_sdk_gc_cleanup() -> None: + client = create_feature_usage_http_client() + close = AsyncMock() + + with patch.object(client, "aclose", close): + cast(Any, client).__del__() + await asyncio.sleep(0) + + close.assert_awaited_once() + await client.aclose() + + async def test_ensure_async_token_provider_wraps_sync_provider() -> None: def sync_provider() -> str: return "sync-token" From fcb8f57ff7dc589fac40d1a3570076a7619a360b Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 30 Jul 2026 09:04:34 +0200 Subject: [PATCH 7/8] Python: preserve caller-owned Foundry transports Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- .../packages/foundry/agent_framework_foundry/_agent.py | 7 +++---- .../foundry/agent_framework_foundry/_chat_client.py | 6 ++++-- .../foundry/agent_framework_foundry/_foundry_evals.py | 4 ++-- .../packages/foundry/tests/foundry/test_foundry_agent.py | 9 ++++++--- .../foundry/tests/foundry/test_foundry_chat_client.py | 5 ++++- python/packages/foundry/tests/test_foundry_evals.py | 2 +- python/packages/tools/tests/test_docker_shell_tool.py | 7 ++++++- python/packages/tools/tests/test_local_shell_tool.py | 6 +++++- 8 files changed, 31 insertions(+), 15 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 61b892da06a..2013d0401af 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -269,7 +269,8 @@ def __init__( openai_client_kwargs: dict[str, Any] = {} if default_headers: openai_client_kwargs["default_headers"] = dict(default_headers) - openai_client_kwargs["http_client"] = create_foundry_feature_usage_http_client() + if self._should_close_client: + openai_client_kwargs["http_client"] = create_foundry_feature_usage_http_client() if allow_preview: openai_client_kwargs["agent_name"] = self.agent_name openai_client = self.project_client.get_openai_client(**openai_client_kwargs) @@ -808,9 +809,7 @@ async def create_conversation(self, *, session_id: str | None = None) -> AgentSe Foundry conversation ID. """ client = cast(RawFoundryAgentChatClient, self.client) - conversation = await client.project_client.get_openai_client( - http_client=create_foundry_feature_usage_http_client() - ).conversations.create() + conversation = await client.client.conversations.create() return self.get_session(service_session_id=conversation.id, session_id=session_id) @override diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 6082aecb54b..67f9ba595df 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -209,12 +209,13 @@ def __init__( project_endpoint = foundry_settings.get("project_endpoint") + owns_project_client = project_client is None if project_endpoint is None and project_client is None: raise ValueError( "Either 'project_endpoint' or 'project_client' is required. " "Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable." ) - if not project_client: + if project_client is None: if not project_endpoint: raise ValueError( "Azure AI project endpoint is required. Set via 'project_endpoint' parameter " @@ -237,7 +238,8 @@ def __init__( openai_kwargs: dict[str, Any] = {} if default_headers: openai_kwargs["default_headers"] = default_headers - openai_kwargs["http_client"] = create_foundry_feature_usage_http_client() + if owns_project_client: + openai_kwargs["http_client"] = create_foundry_feature_usage_http_client() super().__init__( model=resolved_model, diff --git a/python/packages/foundry/agent_framework_foundry/_foundry_evals.py b/python/packages/foundry/agent_framework_foundry/_foundry_evals.py index 16e989b7b02..25fd3bd817a 100644 --- a/python/packages/foundry/agent_framework_foundry/_foundry_evals.py +++ b/python/packages/foundry/agent_framework_foundry/_foundry_evals.py @@ -47,7 +47,7 @@ from openai import AsyncOpenAI from ._chat_client import FoundryChatClient -from ._feature_usage import FeatureIndex, create_foundry_feature_usage_http_client +from ._feature_usage import FeatureIndex if TYPE_CHECKING: from azure.ai.projects.aio import AIProjectClient @@ -667,7 +667,7 @@ def _resolve_openai_client( return client.client return client if project_client is not None: - oai = project_client.get_openai_client(http_client=create_foundry_feature_usage_http_client()) + oai = project_client.get_openai_client() if oai is None: # pyright: ignore[reportUnnecessaryComparison] raise ValueError("project_client.get_openai_client() returned None. Check project configuration.") if not isinstance(oai, AsyncOpenAI): diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index c8473fd3689..e6fc1cde136 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -123,7 +123,7 @@ def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None: assert client.agent_name == "test-agent" assert client.agent_version == "1.0" - mock_project.get_openai_client.assert_called_once_with(http_client=ANY) + mock_project.get_openai_client.assert_called_once_with() def test_raw_foundry_agent_chat_client_creates_project_client_with_feature_policy() -> None: @@ -141,6 +141,7 @@ def test_raw_foundry_agent_chat_client_creates_project_client_with_feature_polic assert len(policies) == 1 assert isinstance(policies[0], FeatureUsagePolicy) assert "custom_hook_policy" not in factory.call_args.kwargs + mock_project.get_openai_client.assert_called_once_with(http_client=ANY) async def test_foundry_agent_basic_call_does_not_request_unsupported_encrypted_reasoning() -> None: @@ -238,7 +239,6 @@ def test_raw_foundry_agent_chat_client_init_passes_agent_name_when_preview_enabl mock_project.get_openai_client.assert_called_once_with( agent_name="hosted-agent", default_headers={"x-test": "1"}, - http_client=ANY, ) @@ -1154,12 +1154,13 @@ async def test_foundry_agent_create_conversation_returns_agent_session() -> None mock_project = MagicMock() mock_project.get_openai_client.return_value = openai_client agent = FoundryAgent(project_client=mock_project, agent_name="test-agent") + mock_project.get_openai_client.reset_mock() session = await agent.create_conversation() assert isinstance(session, AgentSession) assert session.service_session_id == "conv_123" - mock_project.get_openai_client.assert_called() + mock_project.get_openai_client.assert_not_called() openai_client.conversations.create.assert_awaited_once_with() @@ -1171,11 +1172,13 @@ async def test_foundry_agent_create_conversation_accepts_local_session_id() -> N mock_project = MagicMock() mock_project.get_openai_client.return_value = openai_client agent = FoundryAgent(project_client=mock_project, agent_name="test-agent") + mock_project.get_openai_client.reset_mock() session = await agent.create_conversation(session_id="local-session") assert session.session_id == "local-session" assert session.service_session_id == "conv_123" + mock_project.get_openai_client.assert_not_called() def test_foundry_agent_init() -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 128fbb0e04d..c8fc8ad154d 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -9,7 +9,7 @@ from functools import wraps from pathlib import Path from typing import Annotated, Any, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import agent_framework._telemetry as telemetry import pytest @@ -230,6 +230,7 @@ def test_init() -> None: assert client.model == _TEST_FOUNDRY_MODEL assert client.project_client is mock_project_client assert isinstance(client, SupportsChatGetResponse) + mock_project_client.get_openai_client.assert_called_once_with() def test_raw_foundry_chat_client_init_uses_explicit_parameters() -> None: @@ -269,6 +270,7 @@ def test_init_with_default_header() -> None: assert client.default_headers is not None assert key in client.default_headers assert client.default_headers[key] == value + project_client.get_openai_client.assert_called_once_with(default_headers=default_headers) def test_init_with_project_endpoint_creates_project_client() -> None: @@ -295,6 +297,7 @@ def test_init_with_project_endpoint_creates_project_client() -> None: assert len(policies) == 1 assert isinstance(policies[0], FeatureUsagePolicy) assert "custom_hook_policy" not in factory.call_args.kwargs + project_client.get_openai_client.assert_called_once_with(http_client=ANY) def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/python/packages/foundry/tests/test_foundry_evals.py b/python/packages/foundry/tests/test_foundry_evals.py index f59772df60a..575a5976bbf 100644 --- a/python/packages/foundry/tests/test_foundry_evals.py +++ b/python/packages/foundry/tests/test_foundry_evals.py @@ -923,7 +923,7 @@ def test_constructor_with_project_client(self) -> None: mock_project.get_openai_client.return_value = mock_oai fe = FoundryEvals(project_client=mock_project, model="gpt-4o") assert fe.name == "Microsoft Foundry" - mock_project.get_openai_client.assert_called_once() + mock_project.get_openai_client.assert_called_once_with() def test_constructor_no_client_auto_creates_from_env(self) -> None: """When no client/project_client given, auto-creates FoundryChatClient from env.""" diff --git a/python/packages/tools/tests/test_docker_shell_tool.py b/python/packages/tools/tests/test_docker_shell_tool.py index da486f8cb31..e9f63fa4070 100644 --- a/python/packages/tools/tests/test_docker_shell_tool.py +++ b/python/packages/tools/tests/test_docker_shell_tool.py @@ -19,6 +19,7 @@ import pytest +from agent_framework_tools._feature_usage import FeatureIndex from agent_framework_tools.shell import ( DockerNotAvailableError, DockerShellTool, @@ -367,10 +368,14 @@ async def test_run_dispatches_to_private_stateless_runner() -> None: tool = DockerShellTool(mode="stateless") expected = ShellResult(stdout="ok", stderr="", exit_code=0, duration_ms=1) - with patch.object(tool, "_run_stateless", AsyncMock(return_value=expected)) as run_stateless: + with ( + patch.object(tool, "_run_stateless", AsyncMock(return_value=expected)) as run_stateless, + patch("agent_framework_tools.shell._docker.mark_feature_used") as mark_feature_used, + ): result = await tool.run("echo hi", timeout=9.0) assert result is expected + mark_feature_used.assert_called_once_with(FeatureIndex.TOOLS_SHELL) run_stateless.assert_awaited_once_with("echo hi", timeout=9.0) diff --git a/python/packages/tools/tests/test_local_shell_tool.py b/python/packages/tools/tests/test_local_shell_tool.py index 5f5488bb296..d47935c1f75 100644 --- a/python/packages/tools/tests/test_local_shell_tool.py +++ b/python/packages/tools/tests/test_local_shell_tool.py @@ -7,6 +7,7 @@ import pytest +from agent_framework_tools._feature_usage import FeatureIndex from agent_framework_tools.shell import LocalShellTool, ShellCommandError, ShellPolicy from agent_framework_tools.shell._executor import _popen_kwargs_for_group, run_stateless @@ -34,7 +35,10 @@ async def communicate(self) -> tuple[bytes, bytes]: async def test_stateless_echo() -> None: tool = LocalShellTool(mode="stateless", approval_mode="never_require", acknowledge_unsafe=True) cmd = "Write-Output hello" if sys.platform == "win32" else "echo hello" - result = await tool.run(cmd) + with patch("agent_framework_tools.shell._tool.mark_feature_used") as mark_feature_used: + result = await tool.run(cmd) + + mark_feature_used.assert_called_once_with(FeatureIndex.TOOLS_SHELL) assert "hello" in result.stdout assert result.exit_code == 0 assert result.timed_out is False From acdafd62727eb717a87318cee0fb0dceafd592fb Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 30 Jul 2026 12:16:52 +0200 Subject: [PATCH 8/8] Python: remove stale Anthropic test import Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 346bf168-b668-4c4a-a8db-a67282ee5e5f --- python/packages/anthropic/tests/test_anthropic_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index d713111152e..342302d4d22 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -16,7 +16,6 @@ FunctionInvocationLayer, Message, SupportsChatGetResponse, - UsageDetails, tool, ) from agent_framework._settings import load_settings