diff --git a/README.md b/README.md index c6f4e47..fbf8f51 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,18 @@ If `AZURE_FUNCTIONS_AGENTS_PROVIDER` is unset, auto-detection picks the first pr Model resolution precedence is: explicit requested model > provider-specific env (`FOUNDRY_MODEL` for Foundry, `AZURE_OPENAI_DEPLOYMENT` for Azure OpenAI) > `AZURE_FUNCTIONS_AGENTS_MODEL` > provider default. +### Provider client lifetime + +The runtime reuses provider SDK clients and their HTTP connection pools for the +life of each Python worker process. This applies uniformly to built-in chat +endpoints, non-HTTP triggers, delegated agents, and workflow sub-agent +activities. MAF `Agent` objects are still created per invocation because they +carry mutable request state. + +Provider settings are process configuration. Deployments and app-setting +updates recycle Functions workers, so each new worker builds clients from the +new settings; mutating `os.environ` inside a running worker is unsupported. + ## Quick Start ### 1. Create the agent file diff --git a/docs/architecture.md b/docs/architecture.md index 492114b..b1ff495 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -62,7 +62,7 @@ A few boundaries are worth calling out explicitly: | `azure_functions_agents/system_tools/sandbox.py` | Builds the ACA Dynamic Sessions-backed `execute_python` tool for a resolved agent/session, using a fresh GUID when no explicit session id is provided. | `create_sandbox_tools()` | | `azure_functions_agents/system_tools/web_request.py` | Builds the default-on, SSRF-guarded `web_request` outbound HTTP tool, built once per agent at registration (no Azure resource required). | `create_web_request_tools()` | | `azure_functions_agents/runner.py` | Executes prompts through the Microsoft Agent Framework, managing sessions, tools, and streaming; builds per-request `delegate_` tools and fresh stateless workflow leaf agents; attempts one internal token-usage record through the shared runtime logger for each actual MAF invocation attempt. | `run_agent()`, `run_agent_stream()`, `build_subagent_tools()`, `run_leaf_agent_task()` | -| `azure_functions_agents/client_manager.py` | Defines the pluggable inference-client abstraction, immutable inference-target metadata, and the default MAF-backed implementation. | `ClientManager`, `InferenceTarget`, `get_client_manager()`, `set_client_manager()` | +| `azure_functions_agents/client_manager.py` | Defines the pluggable inference-client abstraction, immutable inference-target metadata, and the default MAF-backed implementation. The default manager owns a worker-process cache of typed provider clients/connection pools keyed by provider and model, plus a shared async credential. | `ClientManager`, `InferenceTarget`, `get_client_manager()`, `set_client_manager()`, `shutdown_client_manager()` | | `azure_functions_agents/workflows/*` | Experimental Dynamic Workflow runtime: Durable orchestration registration, workflow tool and Sub Agent execution, immutable owner policy, plan validation/schema, session ownership, and workflow-management tools. | `register_workflows()`, `build_workflow_integration()`, `WorkflowPlanPolicy` | | `azure_functions_agents/_function_tool.py` | Thin local shim around MAF `FunctionTool` creation so project tools can use `@tool`, plus `@workflow_tool` metadata for Dynamic Workflow Activity targets. | `tool()`, `workflow_tool()` | | `azure_functions_agents/_logger.py` | Shared package logger used across discovery, registration, and runtime code. | `logger` | @@ -168,7 +168,7 @@ The `create_function_app()` docstring in `src/azure_functions_agents/app.py:crea ### Where the registration stage hands off to execution -Registration does not run the agent itself. Instead, `registration/_handlers.py` builds closures that call `runner.run_agent()` or `runner.run_agent_stream()`, passing the `ResolvedAgent` instructions plus the already-filtered `AgentCapabilities` — and, when the agent declares `subagents`, its `ResolvedAgent.subagents` list plus the frozen `AgentCatalog`. For non-HTTP triggers, the closure delegates payload construction to `registration/_trigger_serialization.py`: native `to_dict()`/`model_dump()` contracts are used first, then public Azure Functions binding adapters, batch recursion, and byte encoding produce JSON-safe prompt data. HTTP handlers build their request-body JSON separately and do not use this serializer. The runner then asks the active `ClientManager` to build a chat client, builds any `delegate_` tools fresh for this request, and executes through the Microsoft Agent Framework (`src/azure_functions_agents/runner.py`, `src/azure_functions_agents/client_manager.py`). +Registration does not run the agent itself. Instead, `registration/_handlers.py` builds closures that call `runner.run_agent()` or `runner.run_agent_stream()`, passing the `ResolvedAgent` instructions plus the already-filtered `AgentCapabilities` — and, when the agent declares `subagents`, its `ResolvedAgent.subagents` list plus the frozen `AgentCatalog`. For non-HTTP triggers, the closure delegates payload construction to `registration/_trigger_serialization.py`: native `to_dict()`/`model_dump()` contracts are used first, then public Azure Functions binding adapters, batch recursion, and byte encoding produce JSON-safe prompt data. HTTP handlers build their request-body JSON separately and do not use this serializer. The runner then asks the active `ClientManager` for the provider/model chat client, builds any `delegate_` tools fresh for this request, and executes through the Microsoft Agent Framework (`src/azure_functions_agents/runner.py`, `src/azure_functions_agents/client_manager.py`). The default manager returns a process-cached client for an identical target rather than constructing a new SDK client and connection pool per invocation. For a workflow-enabled main agent, `workflows/integration.py` produces one immutable `WorkflowPlanPolicy` from the concrete workflow tools and @@ -343,14 +343,64 @@ This split keeps parsing, policy, Azure binding registration, and runtime execut ### Custom inference client -To plug in a different chat backend, implement the `ClientManager` interface and register it once with `set_client_manager(...)`; after that, `runner.run_agent()` and `runner.run_agent_stream()` use your implementation for every call. See `src/azure_functions_agents/client_manager.py` and the README section [Plugging in a custom client manager](https://github.com/Azure/azure-functions-agents-runtime/blob/main/README.md#plugging-in-a-custom-client-manager). +To plug in a different chat backend, implement the `ClientManager` interface and register it once with `set_client_manager(...)`; after that, `runner.run_agent()` and `runner.run_agent_stream()` use your implementation for every call. See `src/azure_functions_agents/client_manager.py` and [Provider-client lifetime](#provider-client-lifetime). This extension point is deliberately below the registration layer: no trigger or endpoint code needs to change when you swap providers. The `ResolvedAgent.model` value is still the hand-off contract, but your manager decides how to interpret it. Delegated specialists resolve their model through the same `ClientManager`, so a custom implementation applies uniformly to coordinators and specialists alike. +Install a custom manager before the default manager is first requested. Any +active manager, default or custom, must be closed with +`await shutdown_client_manager()` before `set_client_manager(...)` installs a +replacement. While shutdown is in progress, both lookup and replacement are +rejected so no caller can acquire a closing manager or bypass its cleanup. + The runner calls `build_chat_client_with_target()` and receives the client plus a frozen `InferenceTarget` containing nullable `provider` and `model` fields. Its concrete base implementation calls the existing abstract `build_chat_client()` once and returns an empty descriptor, so existing custom managers remain compatible. A custom manager can override the new method when it can authoritatively describe the target used to construct its client. `MAFClientManager` resolves provider and effective model once for client construction and metadata. Subclasses that override the existing `build_chat_client()` hook retain that dispatch and receive an empty descriptor; they can override `build_chat_client_with_target()` when they can provide authoritative metadata. +### Provider-client lifetime + +`MAFClientManager` is a process-wide owner, not only a factory. It lazily caches +one built-in MAF chat client for each provider and resolved model. Primary agents, declared triggers, +built-in endpoints, delegated specialists, and Workflow Sub Agent activities +therefore reuse the same SDK connection pool when they resolve to the same +target. `Agent`, `AgentSession`, history, and tools remain per run because those +objects carry mutable request state. + +The cached MAF wrapper is shared by concurrent async invocations. Characterization +tests against the pinned Agent Framework version drive real `OpenAIChatClient` +and `FoundryChatClient` wrappers through overlapping streaming, non-streaming, +and mixed calls. They verify that prompts, sessions, function-call IDs, and +stream events remain isolated without serializing requests. An explicit upstream +compatibility contract is tracked in +[microsoft/agent-framework#7654](https://github.com/microsoft/agent-framework/issues/7654). + +This follows the Azure Functions recommendation to +[reuse SDK client instances across invocations](https://learn.microsoft.com/azure/azure-functions/manage-connections#manage-sdk-client-connections). +Each Python language-worker process owns its own manager and cache. Async +invocations in that worker share its persistent event loop, which is also the +lifetime boundary for the cached aiohttp/httpx transports. Provider settings +are inherited from the worker process environment; deployments and app-setting +updates recycle workers rather than mutating a running process environment. + +Managed-identity-backed Azure OpenAI and Foundry clients share one +manager-owned async credential and token cache. `MAFClientManager.close()` is +idempotent and closes cached client transports before that credential. +Agent Framework 1.3 does not expose a public close method on its chat clients, +so an internal typed ownership wrapper records the pinned implementation's +`client` (`AsyncOpenAI`/httpx) and, for Foundry, the independent +`project_client` (`AIProjectClient`/aiohttp) when each chat client is built. +Contract tests guard this version-sensitive integration and fail fast if those +attributes change. + +Azure Functions' Python `FunctionApp` has no supported async shutdown callback. +The runtime therefore does not install `atexit` or process-signal handlers that +could run cleanup on the wrong or an already-closed event loop. Strong +process-lifetime ownership prevents the unbounded per-invocation session leak +and connection churn reported by Issue #157; `shutdown_client_manager()` is the +explicit cleanup API for tests and embedding hosts that own an async lifecycle. +The missing host lifecycle contract is tracked in +[azure-functions-python-worker#1904](https://github.com/Azure/azure-functions-python-worker/issues/1904). + ### Custom tools To add project-specific tools, drop a `.py` file into `tools/` and expose either `@tool`-decorated functions or plain functions that can be auto-wrapped into `FunctionTool` objects. Discovery lives in `src/azure_functions_agents/discovery/tools.py:discover_project_tools()` (with `discover_user_tools()` kept as the normal-tool compatibility API), and the local decorator shim is in `src/azure_functions_agents/_function_tool.py:tool()`. diff --git a/docs/frds/0007-multi-agent-delegation.md b/docs/frds/0007-multi-agent-delegation.md index 86325d7..0581f12 100644 --- a/docs/frds/0007-multi-agent-delegation.md +++ b/docs/frds/0007-multi-agent-delegation.md @@ -381,8 +381,9 @@ Do not build specialists once at startup: confirmed in pinned `agent-framework-core==1.3.*` and remains true upstream. One warm Functions worker can serve concurrent requests, so sharing one live agent would create the race class already guarded by the per-session lock. -- Fresh construction is cheap. `ClientManager` already reuses the expensive - model client process-wide. A lightweight `Agent` wrapper adds little work, +- Fresh construction is cheap. `ClientManager` reuses the expensive model + client process-wide (the implementation was corrected by Issue #157; see + Decision #21). A lightweight `Agent` wrapper adds little work, prevents cross-request state sharing, and uses less cold-start time and memory than pre-building every agent. @@ -439,8 +440,8 @@ individually expanded remote functions). Construct the specialist's `FunctionTool` wrapper while assembling one invocation's coordinator tools; the specialist `Agent` itself is not built there — its handler builds one fresh on every call (§4.7, §5 Decision #20). -`ClientManager` builds each specialist client for that specialist's own -resolved model and reuses provider/credential state process-wide. Do not +`ClientManager` resolves each specialist client for that specialist's own +model and reuses the matching provider client and credential process-wide. Do not cache mutable MAF agents between Functions requests, or between calls within one request. Declaring a specialist creates one cheap wrapper and adds its schema to the coordinator prompt; the specialist's model does not run until @@ -639,6 +640,10 @@ handoff participant may itself declare `subagents` and delegate. | 18 | Who may declare `subagents` | any independently runnable agent / main-agent-only (mirror FRD 0004 `workflows.enabled`) | **Any independently runnable agent** may declare `subagents`. Single-level (#6) still applies, so when that agent is itself invoked as a sub-agent its `subagents` are not wired and it cannot delegate onward. Simpler than a main/non-main split and matches the cross-framework norm (#15) | Human (user) | 2026-07-15 | | 19 | Observability approach | new bespoke tracing / rely on existing auto-instrumentation + add delegation enrichment / defer all enrichment | **Rely on auto-instrumentation, add delegation enrichment** — the runtime already enables MAF `gen_ai` spans and Azure Monitor export (`_observability.py`), and the delegate tool calling `Agent.run()` (originally via `as_tool()`→`run()`; the hand-written tool added by #20 calls `run()` directly, same effect) + `FunctionTool.invoke()` auto-nest a delegated call under one trace/`OperationId` with no new tracing code (verified against MAF tag `python-1.3.0` and the Functions Python worker's context attach). v1 additionally adds `af.delegate.*` attributes, delegate metrics, and explicit delegated-error accounting for parity with sandbox/web_request. Token roll-up across the boundary and SSE stream-through of specialist internals are documented limitations | Human (user) | 2026-07-15 | | 20 | Delegate execution mechanism, revisited post-implementation | keep `as_tool()` + per-specialist `asyncio.Lock` + `specialist_agent.run` monkeypatch/stream-capture (as first implemented, #1/#3) / rewrite as a hand-written non-streaming `@tool(schema=...)` function tool, building a fresh specialist `Agent` per call | **Hand-written non-streaming tool, built fresh per call** — a delegate only ever needs the specialist's final text (§4.12's "SSE is a black box at the boundary" was already a non-goal), so there is no reason to run the specialist through `Agent.run(stream=True, ...)` at all, which is all `as_tool()`'s own `_agent_wrapper` ever did internally before `await`-ing `stream.get_final_response()` back into one string anyway. That streaming requirement was the *only* reason the first implementation needed to monkeypatch `specialist_agent.run` (to capture the `ResponseStream` `as_tool()` builds internally, so it could be force-finalized on timeout/cancellation — `ResponseStream.__anext__`'s cleanup hooks only fire from its own `except StopAsyncIteration`/`except Exception` branches, never `BaseException`) and, because the specialist `Agent` object was shared across calls in a turn, to serialize concurrent same-specialist calls behind a per-specialist `asyncio.Lock` so that monkeypatch rebind was race-free. Switching to plain, non-streaming `agent.run(task)` — verified against installed `agent-framework-core==1.3.0` (`AgentTelemetryLayer._run`, `ChatTelemetryLayer._get_response` in `agent_framework.observability`) to close its OTel spans deterministically on *any* exception, `asyncio.CancelledError` included, via the ordinary `with`/context-manager `__exit__` guarantee (no `BaseException` gap like the streaming path has) — removes the need to capture or finalize a stream at all. Building the specialist `Agent` fresh on every call, instead of once per tool-build and reused, removes the shared mutable state the lock existed to protect, so the lock is removed too: same-specialist calls now simply run in parallel, each on its own instance (revises #14). Net result: less code, a cleaner and more debuggable per-call span timeline (each specialist run's span opens/closes at one well-defined point per call, on every path — success, recoverable failure, timeout, cancel), and no behavior change visible to the coordinator or its model | Human (user) | 2026-07-16 | +| 21 | Provider-client lifecycle after Issue #157 exposed per-run Foundry session leaks | construct and close per run / cache only Foundry / cache all built-in MAF provider clients per worker process | **Cache built-in provider clients by provider, resolved model, and non-secret endpoint configuration in the process-wide `MAFClientManager`; share its managed-identity credential and keep `Agent` instances per call.** This implements the process-wide reuse already required by §4.7, follows Azure Functions SDK-client guidance, and avoids per-run TCP/TLS churn and SNAT pressure. The Functions Python worker runs async invocations on one persistent worker event loop; cached async transports share that loop. Cleanup remains explicit because `FunctionApp` has no supported async shutdown callback. With pinned Agent Framework 1.3, full Foundry cleanup requires awaiting both private attributes (`chat_client.client.close()` for httpx and `chat_client.project_client.close()` for aiohttp); contract tests and warnings guard that version coupling. | Human (user), Issue #157 | 2026-08-12 | +| 22 | Simplify provider-client ownership after implementation review | detached cleanup task plus default-manager retirement / direct awaited cleanup plus one singleton shutdown state | **Use direct awaited cleanup and one process-singleton shutdown state.** Any active manager, including a custom one, requires explicit shutdown before replacement; lookup and replacement fail while shutdown runs. Cache entries are keyed by provider and resolved model, while endpoint/API/auth configuration is fixed after that provider's first build and a change requires a worker restart. This preserves process-wide connection reuse without detached tasks or manager-specific replacement rules. | Human (user), Laveesh review | 2026-08-12 | +| 23 | Type and configuration handling after second implementation review | dynamic `Any`/`getattr` cleanup and runtime config-change rejection / typed ownership wrapper and process-environment configuration | **Use a typed internal ownership wrapper and treat provider settings as worker process configuration.** Each cache entry records its concrete MAF chat client and typed closeable transports at construction, so pinned Agent Framework contract drift fails immediately instead of being discovered through cleanup-time `getattr`. The shared async credential is typed as well. Remove runtime endpoint/API/auth mutation tracking: Functions deployment and app-setting changes recycle workers, while mutating `os.environ` inside a live worker is unsupported. Keep the custom-manager replacement note narrowly scoped to test, custom provider, and embedding-host scenarios. | Human (user), Laveesh review | 2026-08-12 | +| 24 | Concurrent use of one cached MAF wrapper after streaming review | serialize all runs / create wrappers per run over one transport / validate and share the pinned wrappers | **Validate and share the pinned wrappers without a global lock.** Deterministic tests drive real `OpenAIChatClient` and `FoundryChatClient` wrappers through overlapping streaming + streaming, non-streaming + non-streaming, and mixed calls. They verify distinct prompts, sessions, function-call IDs, and stream events do not cross over and that requests overlap rather than serialize. Fresh `Agent` and `AgentSession` instances remain the mutable per-run boundary. A lock would hold for an entire stream and unnecessarily destroy throughput; constructing wrappers per run would rely on private transport injection. The missing formal upstream concurrent-use guarantee is tracked in microsoft/agent-framework#7654. | Human (user), Laveesh streaming review | 2026-08-12 | ## 6. Test plan @@ -672,6 +677,14 @@ handoff participant may itself declare `subagents` and delegate. build and run on their own independent instance, in parallel, both producing correct results (no shared-instance lock — #20); calls to different specialists also run in parallel; verify every result. +- [ ] Regression: repeated primary/delegated builds for one provider/model reuse + one process-owned MAF client and credential; distinct model/endpoint keys + remain isolated; explicit manager shutdown awaits both Foundry transports + exactly once and continues cleanup after an individual close failure + (Issue #157, Decision #21). Concurrent real-wrapper runs cover + streaming + streaming, non-streaming + non-streaming, and mixed calls, + proving request/session/function-call/event isolation without serialization + (Decision #24). - [ ] Observability — with instrumentation enabled, one delegated call produces nested `execute_tool delegate_` and `invoke_agent {specialist}` spans under the coordinator's `agent.run` span, all sharing one trace id; diff --git a/src/azure_functions_agents/client_manager.py b/src/azure_functions_agents/client_manager.py index 254f335..357ffa4 100644 --- a/src/azure_functions_agents/client_manager.py +++ b/src/azure_functions_agents/client_manager.py @@ -7,33 +7,40 @@ Only one implementation ships today: :class:`MAFClientManager`. It is selected automatically by :func:`get_client_manager` and lives behind a process-wide -singleton because building a provider client (and the underlying credential -caches it owns) is cheap to share across requests. +singleton. It caches provider clients by resolved target so one Python worker +reuses the underlying HTTP connection pools across requests. ABC surface ----------- * :meth:`ClientManager.resolve_model` — pick the actual model/deployment to use given an optional per-call request. -* :meth:`ClientManager.build_chat_client` — return a fresh ``ChatClient`` - bound to a specific model. -* :meth:`ClientManager.build_chat_client_with_target` — return a fresh client - with authoritative inference-target metadata when available. +* :meth:`ClientManager.build_chat_client` — return a ``ChatClient`` bound to a + specific model. +* :meth:`ClientManager.build_chat_client_with_target` — return a client with + authoritative inference-target metadata when available. * :meth:`ClientManager.close` — release any resources held by the manager - (called from the application's shutdown hook). + when an embedding host owns an async shutdown lifecycle. """ from __future__ import annotations import os +import threading from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any, Protocol from ._credential import build_async_credential from ._logger import logger from .config.env import runtime_env_value +if TYPE_CHECKING: + from agent_framework import BaseChatClient + from agent_framework.foundry import FoundryChatClient + from agent_framework.openai import OpenAIChatClient + from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential + # --------------------------------------------------------------------------- # ABC # --------------------------------------------------------------------------- @@ -47,6 +54,44 @@ class InferenceTarget: model: str | None = None +@dataclass(frozen=True) +class _ClientCacheKey: + provider: str + model: str + + +class _AsyncCloseable(Protocol): + async def close(self) -> None: ... + + +@dataclass(frozen=True) +class _OwnedResource: + value: _AsyncCloseable + label: str + + +@dataclass(frozen=True) +class _ManagedChatClient: + client: BaseChatClient[Any] + resources: tuple[_OwnedResource, ...] + + async def close(self) -> list[Exception]: + return await _close_owned_resources(self.resources) + + +async def _close_owned_resources( + resources: tuple[_OwnedResource, ...], +) -> list[Exception]: + errors: list[Exception] = [] + for resource in resources: + try: + await resource.value.close() + except Exception as exc: + logger.error("Failed to close %s: %s", resource.label, exc) + errors.append(exc) + return errors + + class ClientManager(ABC): """Provider-agnostic interface for building chat clients.""" @@ -103,6 +148,12 @@ class MAFClientManager(ClientManager): name = "maf" + def __init__(self) -> None: + self._clients: dict[_ClientCacheKey, _ManagedChatClient] = {} + self._async_credential: AsyncDefaultAzureCredential | None = None + self._lock = threading.RLock() + self._closed = False + def resolve_model(self, requested: str | None) -> str: """Resolve model as requested > provider-specific env > runtime env > default.""" return self._resolve_model(requested, self._provider()) @@ -140,19 +191,17 @@ def _build_maf_chat_client_with_target( ) -> tuple[Any, InferenceTarget]: provider = self._provider() resolved = self._resolve_model(model, provider) - logger.info("MAF provider=%s model=%s", provider, resolved) - if provider == "openai": - client = self._build_openai(resolved) - elif provider == "azure_openai": - client = self._build_azure_openai(resolved) - elif provider == "foundry": - client = self._build_foundry(resolved) - else: - raise RuntimeError( - f"Unknown AZURE_FUNCTIONS_AGENTS_PROVIDER '{provider}'. " - "Use one of: openai, azure_openai, foundry." - ) - return client, InferenceTarget( + cache_key = _ClientCacheKey(provider, resolved) + with self._lock: + self._ensure_open() + managed = self._clients.get(cache_key) + if managed is None: + managed = self._build_provider_client(provider, resolved) + self._clients[cache_key] = managed + logger.info("Created MAF provider client: provider=%s model=%s", provider, resolved) + else: + logger.debug("Reusing MAF provider client: provider=%s model=%s", provider, resolved) + return managed.client, InferenceTarget( provider=provider, model=resolved, ) @@ -161,6 +210,36 @@ def _build_maf_chat_client_with_target( # Internals # ------------------------------------------------------------------ + def _build_provider_client(self, provider: str, model: str) -> _ManagedChatClient: + if provider == "openai": + client = self._build_openai(model) + return _ManagedChatClient( + client, + (_OwnedResource(client.client, "OpenAI AsyncOpenAI transport"),), + ) + if provider == "azure_openai": + client = self._build_azure_openai(model) + return _ManagedChatClient( + client, + (_OwnedResource(client.client, "Azure OpenAI AsyncOpenAI transport"),), + ) + if provider == "foundry": + client = self._build_foundry(model) + return _ManagedChatClient( + client, + ( + _OwnedResource(client.client, "Foundry AsyncOpenAI transport"), + _OwnedResource( + client.project_client, + "Foundry AIProjectClient transport", + ), + ), + ) + raise RuntimeError( + f"Unknown AZURE_FUNCTIONS_AGENTS_PROVIDER '{provider}'. " + "Use one of: openai, azure_openai, foundry." + ) + @staticmethod def _env(name: str) -> str: """Return ``$name`` stripped, or ``""`` if missing/blank. @@ -192,7 +271,7 @@ def _provider(cls) -> str: ) @classmethod - def _build_openai(cls, model: str) -> Any: + def _build_openai(cls, model: str) -> OpenAIChatClient: from agent_framework.openai import OpenAIChatClient return OpenAIChatClient( @@ -200,11 +279,10 @@ def _build_openai(cls, model: str) -> Any: api_key=cls._env("OPENAI_API_KEY") or None, ) - @classmethod - def _build_azure_openai(cls, model: str) -> Any: + def _build_azure_openai(self, model: str) -> OpenAIChatClient: from agent_framework.openai import OpenAIChatClient - endpoint = cls._env("AZURE_OPENAI_ENDPOINT") + endpoint = self._env("AZURE_OPENAI_ENDPOINT") if not endpoint: raise RuntimeError( "AZURE_FUNCTIONS_AGENTS_PROVIDER=azure_openai requires " @@ -217,21 +295,20 @@ def _build_azure_openai(cls, model: str) -> Any: # Only forward api_version when the user explicitly sets it. MAF defaults # to the Responses API ("preview") which rejects Chat Completions GA # versions like "2024-10-21" with "API version not supported". - api_version = cls._env("AZURE_OPENAI_API_VERSION") + api_version = self._env("AZURE_OPENAI_API_VERSION") if api_version: kwargs["api_version"] = api_version - api_key = cls._env("AZURE_OPENAI_API_KEY") + api_key = self._env("AZURE_OPENAI_API_KEY") if api_key: kwargs["api_key"] = api_key else: - kwargs["credential"] = build_async_credential() + kwargs["credential"] = self._get_async_credential() return OpenAIChatClient(**kwargs) - @classmethod - def _build_foundry(cls, model: str) -> Any: + def _build_foundry(self, model: str) -> FoundryChatClient: from agent_framework.foundry import FoundryChatClient - endpoint = cls._env("FOUNDRY_PROJECT_ENDPOINT") + endpoint = self._env("FOUNDRY_PROJECT_ENDPOINT") if not endpoint: raise RuntimeError( "AZURE_FUNCTIONS_AGENTS_PROVIDER=foundry requires " @@ -240,15 +317,51 @@ def _build_foundry(cls, model: str) -> Any: return FoundryChatClient( project_endpoint=endpoint, model=model, - credential=build_async_credential(), + credential=self._get_async_credential(), ) + def _get_async_credential(self) -> AsyncDefaultAzureCredential: + with self._lock: + self._ensure_open() + if self._async_credential is None: + self._async_credential = build_async_credential() + return self._async_credential + + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("MAFClientManager is closed and cannot create resources.") + + async def close(self) -> None: + """Close all provider transports and the shared credential exactly once.""" + with self._lock: + if self._closed: + return + self._closed = True + cached_clients = list(self._clients.values()) + self._clients.clear() + credential = self._async_credential + self._async_credential = None + + errors: list[Exception] = [] + for chat_client in cached_clients: + errors.extend(await chat_client.close()) + if credential is not None: + errors.extend( + await _close_owned_resources( + (_OwnedResource(credential, "shared async credential"),) + ) + ) + if errors: + raise ExceptionGroup("Failed to close one or more MAF client resources.", errors) + # --------------------------------------------------------------------------- # Process-wide singleton selection # --------------------------------------------------------------------------- -_INSTANCE: ClientManager | None = None +_SHUTTING_DOWN = object() +_INSTANCE: ClientManager | object | None = None +_INSTANCE_LOCK = threading.Lock() def get_client_manager() -> ClientManager: @@ -259,27 +372,52 @@ def get_client_manager() -> ClientManager: between alternative implementations. """ global _INSTANCE - if _INSTANCE is None: - _INSTANCE = MAFClientManager() - logger.info("ClientManager initialized: %s", _INSTANCE.name) - return _INSTANCE + with _INSTANCE_LOCK: + if _INSTANCE is _SHUTTING_DOWN: + raise RuntimeError("Client manager shutdown is in progress") + if _INSTANCE is None: + _INSTANCE = MAFClientManager() + logger.info("ClientManager initialized: %s", _INSTANCE.name) + assert isinstance(_INSTANCE, ClientManager) + return _INSTANCE def set_client_manager(manager: ClientManager) -> None: """Override the process-wide :class:`ClientManager`. Intended for tests and for advanced apps that want to plug in a custom - backend. + backend. Call and await :func:`shutdown_client_manager` before replacing + any active manager, including a custom implementation. """ global _INSTANCE - _INSTANCE = manager + with _INSTANCE_LOCK: + if _INSTANCE is _SHUTTING_DOWN: + raise RuntimeError("Client manager shutdown is in progress") + current = _INSTANCE + if current is manager: + return + if current is not None: + raise RuntimeError( + "Cannot replace the active ClientManager. " + "Await shutdown_client_manager() before replacing it." + ) + _INSTANCE = manager async def shutdown_client_manager() -> None: - """Close the active manager (if any). Idempotent.""" + """Close the active manager; sequential calls are safe, concurrent calls are rejected.""" global _INSTANCE - if _INSTANCE is not None: - try: - await _INSTANCE.close() - finally: - _INSTANCE = None + with _INSTANCE_LOCK: + manager = _INSTANCE + if manager is _SHUTTING_DOWN: + raise RuntimeError("Client manager shutdown is already in progress") + if manager is None: + return + _INSTANCE = _SHUTTING_DOWN + assert isinstance(manager, ClientManager) + try: + await manager.close() + finally: + with _INSTANCE_LOCK: + if _INSTANCE is _SHUTTING_DOWN: + _INSTANCE = None diff --git a/tests/test_client_manager.py b/tests/test_client_manager.py index 22d3f7e..3cbc68d 100644 --- a/tests/test_client_manager.py +++ b/tests/test_client_manager.py @@ -1,9 +1,14 @@ from __future__ import annotations +import asyncio +import time +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace from typing import Any -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest +import pytest_asyncio from azure_functions_agents._credential import build_async_credential from azure_functions_agents.client_manager import ( @@ -12,9 +17,29 @@ ClientManager, InferenceTarget, MAFClientManager, + get_client_manager, + set_client_manager, + shutdown_client_manager, ) +@pytest_asyncio.fixture(autouse=True) +async def _reset_process_client_manager() -> None: + await shutdown_client_manager() + yield + await shutdown_client_manager() + + +def _fake_provider_client(provider: str) -> SimpleNamespace: + client = SimpleNamespace(close=AsyncMock()) + if provider == "foundry": + return SimpleNamespace( + client=client, + project_client=SimpleNamespace(close=AsyncMock()), + ) + return SimpleNamespace(client=client) + + @pytest.mark.parametrize( ("provider", "provider_env", "provider_model"), [ @@ -125,7 +150,7 @@ def test_build_chat_client_with_target_matches_client_branch( monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", provider) if endpoint_name and endpoint: monkeypatch.setenv(endpoint_name, endpoint) - client = object() + client = _fake_provider_client(provider) with patch.object(MAFClientManager, builder, return_value=client) as build: built_client, target = MAFClientManager().build_chat_client_with_target("model-one") @@ -140,7 +165,7 @@ def test_maf_target_uses_one_provider_and_model_resolution_pass( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") - client = object() + client = _fake_provider_client("foundry") with ( patch.object(MAFClientManager, "_provider", return_value="foundry") as provider, @@ -158,6 +183,442 @@ def test_maf_target_uses_one_provider_and_model_resolution_pass( resolve.assert_called_once_with("requested-model", "foundry") +@pytest.mark.asyncio +async def test_maf_manager_reuses_provider_client_on_one_worker_loop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + client = _fake_provider_client("foundry") + + with patch.object(MAFClientManager, "_build_foundry", return_value=client) as build: + first, first_target = manager.build_chat_client_with_target("shared-model") + await asyncio.sleep(0) + second, second_target = manager.build_chat_client_with_target("shared-model") + + assert first is client + assert second is first + assert first_target == second_target == InferenceTarget("foundry", "shared-model") + build.assert_called_once_with("shared-model") + + +def test_maf_manager_logs_cache_creation_at_info_and_hit_at_debug( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") + manager = MAFClientManager() + caplog.set_level("DEBUG", logger="azure.functions.AgentRuntime") + + with patch.object( + MAFClientManager, + "_build_openai", + return_value=_fake_provider_client("openai"), + ): + manager.build_chat_client_with_target("shared-model") + manager.build_chat_client_with_target("shared-model") + + created = [record for record in caplog.records if "Created MAF provider client" in record.message] + reused = [record for record in caplog.records if "Reusing MAF provider client" in record.message] + assert len(created) == 1 + assert created[0].levelname == "INFO" + assert len(reused) == 1 + assert reused[0].levelname == "DEBUG" + + +@pytest.mark.parametrize( + ("provider", "endpoint_name", "endpoint", "builder"), + [ + ( + "azure_openai", + "AZURE_OPENAI_ENDPOINT", + "https://account.openai.azure.com", + "_build_azure_openai", + ), + ( + "foundry", + "FOUNDRY_PROJECT_ENDPOINT", + "https://project.example", + "_build_foundry", + ), + ], +) +def test_maf_manager_reuses_auto_detected_provider_client( + monkeypatch: pytest.MonkeyPatch, + provider: str, + endpoint_name: str, + endpoint: str, + builder: str, +) -> None: + monkeypatch.delenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", raising=False) + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False) + monkeypatch.delenv("FOUNDRY_PROJECT_ENDPOINT", raising=False) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv(endpoint_name, endpoint) + manager = MAFClientManager() + client = _fake_provider_client(provider) + + with patch.object(MAFClientManager, builder, return_value=client) as build: + first, first_target = manager.build_chat_client_with_target("shared-model") + second, second_target = manager.build_chat_client_with_target("shared-model") + + assert first is client + assert second is first + assert first_target == second_target == InferenceTarget(provider, "shared-model") + build.assert_called_once_with("shared-model") + + +def test_maf_manager_partitions_cached_clients_by_resolved_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + clients = [ + _fake_provider_client("foundry"), + _fake_provider_client("foundry"), + ] + + with patch.object(MAFClientManager, "_build_foundry", side_effect=clients) as build: + first, _ = manager.build_chat_client_with_target("model-one") + second, _ = manager.build_chat_client_with_target("model-two") + + assert first is clients[0] + assert second is clients[1] + assert first is not second + assert build.call_count == 2 + + +@pytest.mark.parametrize( + ("provider", "builder", "client"), + [ + ("openai", "_build_openai", object()), + ( + "foundry", + "_build_foundry", + SimpleNamespace(client=SimpleNamespace(close=AsyncMock())), + ), + ], +) +def test_maf_manager_fails_fast_when_pinned_client_ownership_contract_is_missing( + monkeypatch: pytest.MonkeyPatch, + provider: str, + builder: str, + client: object, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", provider) + if provider == "foundry": + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + + with ( + patch.object(MAFClientManager, builder, return_value=client), + pytest.raises(AttributeError), + ): + MAFClientManager().build_chat_client_with_target("shared-model") + + +def test_maf_manager_publishes_one_client_during_concurrent_first_use( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + client = _fake_provider_client("foundry") + + def _build(_model: str) -> object: + time.sleep(0.02) + return client + + with ( + patch.object(MAFClientManager, "_build_foundry", side_effect=_build) as build, + ThreadPoolExecutor(max_workers=8) as executor, + ): + futures = [ + executor.submit(manager.build_chat_client_with_target, "shared-model") + for _ in range(8) + ] + built_clients = [future.result()[0] for future in futures] + + assert all(built is client for built in built_clients) + build.assert_called_once_with("shared-model") + + +@pytest.mark.asyncio +async def test_maf_manager_closes_foundry_transports_and_credential_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + openai_client = SimpleNamespace(close=AsyncMock()) + project_client = SimpleNamespace(close=AsyncMock()) + chat_client = SimpleNamespace(client=openai_client, project_client=project_client) + credential = SimpleNamespace(close=AsyncMock()) + + with ( + patch( + "agent_framework.foundry.FoundryChatClient", + return_value=chat_client, + ), + patch( + "azure_functions_agents.client_manager.build_async_credential", + return_value=credential, + ), + ): + built, _ = manager.build_chat_client_with_target("shared-model") + await manager.close() + await manager.close() + + assert built is chat_client + openai_client.close.assert_awaited_once_with() + project_client.close.assert_awaited_once_with() + credential.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_maf_manager_closes_openai_transport_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") + manager = MAFClientManager() + transport = SimpleNamespace(close=AsyncMock()) + chat_client = SimpleNamespace(client=transport) + + with patch.object(MAFClientManager, "_build_openai", return_value=chat_client): + manager.build_chat_client_with_target("shared-model") + await manager.close() + + transport.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_maf_manager_closes_azure_openai_transport_and_credential( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "azure_openai") + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://account.openai.azure.com") + monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False) + manager = MAFClientManager() + transport = SimpleNamespace(close=AsyncMock()) + chat_client = SimpleNamespace(client=transport) + credential = SimpleNamespace(close=AsyncMock()) + + with ( + patch("agent_framework.openai.OpenAIChatClient", return_value=chat_client), + patch( + "azure_functions_agents.client_manager.build_async_credential", + return_value=credential, + ), + ): + manager.build_chat_client_with_target("shared-model") + await manager.close() + + transport.close.assert_awaited_once_with() + credential.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_maf_manager_shares_credential_across_foundry_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + credential = SimpleNamespace(close=AsyncMock()) + chat_clients = [ + SimpleNamespace( + client=SimpleNamespace(close=AsyncMock()), + project_client=SimpleNamespace(close=AsyncMock()), + ), + SimpleNamespace( + client=SimpleNamespace(close=AsyncMock()), + project_client=SimpleNamespace(close=AsyncMock()), + ), + ] + + with ( + patch( + "agent_framework.foundry.FoundryChatClient", + side_effect=chat_clients, + ) as client_ctor, + patch( + "azure_functions_agents.client_manager.build_async_credential", + return_value=credential, + ) as credential_builder, + ): + manager.build_chat_client_with_target("model-one") + manager.build_chat_client_with_target("model-two") + await manager.close() + + credential_builder.assert_called_once_with() + assert [item.kwargs["credential"] for item in client_ctor.call_args_list] == [ + credential, + credential, + ] + + +@pytest.mark.asyncio +async def test_foundry_client_exposes_pinned_transport_cleanup_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv( + "FOUNDRY_PROJECT_ENDPOINT", + "https://example.services.ai.azure.com/api/projects/test", + ) + manager = MAFClientManager() + credential = SimpleNamespace(close=AsyncMock(), get_token=AsyncMock()) + + with patch( + "azure_functions_agents.client_manager.build_async_credential", + return_value=credential, + ): + client, _ = manager.build_chat_client_with_target("test-model") + reused, _ = manager.build_chat_client_with_target("test-model") + + assert reused is client + assert callable(client.client.close) + assert callable(client.project_client.close) + await manager.close() + + credential.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_maf_manager_attempts_all_cleanup_after_close_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + openai_client = SimpleNamespace(close=AsyncMock(side_effect=RuntimeError("httpx close"))) + project_client = SimpleNamespace(close=AsyncMock()) + chat_client = SimpleNamespace(client=openai_client, project_client=project_client) + credential = SimpleNamespace(close=AsyncMock()) + + with ( + patch("agent_framework.foundry.FoundryChatClient", return_value=chat_client), + patch( + "azure_functions_agents.client_manager.build_async_credential", + return_value=credential, + ), + ): + manager.build_chat_client_with_target("shared-model") + with pytest.raises(ExceptionGroup, match="Failed to close"): + await manager.close() + + openai_client.close.assert_awaited_once_with() + project_client.close.assert_awaited_once_with() + credential.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_maf_manager_rejects_build_after_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") + manager = MAFClientManager() + + await manager.close() + + with ( + patch.object(MAFClientManager, "_build_openai", return_value=object()), + pytest.raises(RuntimeError, match="closed"), + ): + manager.build_chat_client_with_target("model-one") + + +def test_maf_manager_close_runs_from_sync_embedding_boundary( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") + manager = MAFClientManager() + transport = SimpleNamespace(close=AsyncMock()) + chat_client = SimpleNamespace(client=transport) + + with patch.object(MAFClientManager, "_build_openai", return_value=chat_client): + manager.build_chat_client_with_target("shared-model") + asyncio.run(manager.close()) + + transport.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_shutdown_blocks_get_and_set_until_cleanup_finishes() -> None: + close_started = asyncio.Event() + allow_close = asyncio.Event() + + class BlockingManager(ClientManager): + def resolve_model(self, requested: str | None) -> str: + return requested or "blocking-model" + + def build_chat_client(self, model: str | None) -> Any: + return object() + + async def close(self) -> None: + close_started.set() + await allow_close.wait() + + blocking = BlockingManager() + replacement = BlockingManager() + set_client_manager(blocking) + shutdown = asyncio.create_task(shutdown_client_manager()) + await close_started.wait() + + with pytest.raises(RuntimeError, match="shutdown is in progress"): + get_client_manager() + with pytest.raises(RuntimeError, match="shutdown is in progress"): + set_client_manager(replacement) + + allow_close.set() + await shutdown + + set_client_manager(replacement) + assert get_client_manager() is replacement + + +@pytest.mark.asyncio +async def test_shutdown_clears_singleton_when_custom_manager_close_fails() -> None: + class FailingManager(ClientManager): + def resolve_model(self, requested: str | None) -> str: + return requested or "failing-model" + + def build_chat_client(self, model: str | None) -> Any: + return object() + + async def close(self) -> None: + raise RuntimeError("close failed") + + set_client_manager(FailingManager()) + + with pytest.raises(RuntimeError, match="close failed"): + await shutdown_client_manager() + + assert isinstance(get_client_manager(), MAFClientManager) + + +@pytest.mark.asyncio +async def test_set_client_manager_requires_shutdown_for_custom_manager() -> None: + class CustomManager(ClientManager): + def resolve_model(self, requested: str | None) -> str: + return requested or "custom-model" + + def build_chat_client(self, model: str | None) -> Any: + return object() + + first = CustomManager() + replacement = CustomManager() + set_client_manager(first) + + with pytest.raises(RuntimeError, match="shutdown_client_manager"): + set_client_manager(replacement) + + await shutdown_client_manager() + set_client_manager(replacement) + assert get_client_manager() is replacement + + def test_custom_manager_target_fallback_builds_client_once() -> None: class CustomManager(ClientManager): calls = 0 @@ -182,7 +643,7 @@ def test_maf_subclass_build_chat_client_override_keeps_virtual_dispatch( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") - provider_client = object() + provider_client = _fake_provider_client("openai") class WrappingMAFClientManager(MAFClientManager): calls = 0 diff --git a/tests/test_client_manager_concurrency.py b/tests/test_client_manager_concurrency.py new file mode 100644 index 0000000..1861267 --- /dev/null +++ b/tests/test_client_manager_concurrency.py @@ -0,0 +1,300 @@ +from __future__ import annotations + +import asyncio +import json +from typing import Any +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +from agent_framework import Agent +from agent_framework.foundry import FoundryChatClient +from agent_framework.openai import OpenAIChatClient +from openai import AsyncOpenAI + +from azure_functions_agents.client_manager import MAFClientManager + +_MARKERS = ("alpha", "beta") + + +def _response(marker: str) -> dict[str, Any]: + return { + "id": f"response-{marker}", + "created_at": 0, + "model": "test-model", + "object": "response", + "output": [ + { + "id": f"message-{marker}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": f"reply-{marker}", + "annotations": [], + "logprobs": [], + } + ], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "status": "completed", + } + + +def _stream_events(marker: str) -> list[dict[str, Any]]: + response = _response(marker) + return [ + { + "type": "response.output_item.added", + "output_index": 0, + "sequence_number": 0, + "item": { + "id": f"function-{marker}", + "type": "function_call", + "call_id": f"call-{marker}", + "name": f"tool_{marker}", + "arguments": "", + "status": "in_progress", + }, + }, + { + "type": "response.function_call_arguments.delta", + "item_id": f"function-{marker}", + "output_index": 0, + "sequence_number": 1, + "delta": json.dumps({"marker": marker}), + }, + { + "type": "response.output_text.delta", + "content_index": 0, + "delta": f"reply-{marker}", + "item_id": f"message-{marker}", + "logprobs": [], + "output_index": 1, + "sequence_number": 2, + }, + { + "type": "response.completed", + "sequence_number": 3, + "response": response, + }, + ] + + +class _DeterministicResponsesTransport: + def __init__(self) -> None: + self.requests: dict[str, list[str]] = {marker: [] for marker in _MARKERS} + self.active_requests = 0 + self.max_active_requests = 0 + + async def __call__(self, request: httpx.Request) -> httpx.Response: + body = request.content.decode() + marker = next(marker for marker in _MARKERS if marker in body) + self.requests[marker].append(body) + self.active_requests += 1 + self.max_active_requests = max(self.max_active_requests, self.active_requests) + try: + await asyncio.sleep(0.01 if marker == "alpha" else 0) + payload = json.loads(body) + if payload.get("stream"): + content = "".join( + f"data: {json.dumps(event)}\n\n" for event in _stream_events(marker) + ) + content += "data: [DONE]\n\n" + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=content, + ) + return httpx.Response(200, json=_response(marker)) + finally: + self.active_requests -= 1 + + +class _FakeProjectClient: + def __init__(self, client: AsyncOpenAI) -> None: + self._client = client + self.close = AsyncMock() + + def get_openai_client(self, **_kwargs: Any) -> AsyncOpenAI: + return self._client + + +def _build_chat_client( + provider: str, + transport: _DeterministicResponsesTransport, +) -> tuple[OpenAIChatClient | FoundryChatClient, _FakeProjectClient | None]: + async_http_client = httpx.AsyncClient(transport=httpx.MockTransport(transport)) + openai_client = AsyncOpenAI( + api_key="test-key", + base_url="https://example.test/v1", + http_client=async_http_client, + ) + function_invocation_configuration = {"enabled": False} + if provider == "foundry": + project_client = _FakeProjectClient(openai_client) + return ( + FoundryChatClient( + project_client=project_client, # type: ignore[arg-type] + model="test-model", + function_invocation_configuration=function_invocation_configuration, + ), + project_client, + ) + return ( + OpenAIChatClient( + model="test-model", + async_client=openai_client, + function_invocation_configuration=function_invocation_configuration, + ), + None, + ) + + +async def _consume_stream( + agent: Agent[Any], + marker: str, + session_id: str, +) -> tuple[str, set[str], str]: + session = agent.create_session(session_id=session_id) + stream = agent.run(marker, stream=True, session=session) + text_parts: list[str] = [] + call_ids: set[str] = set() + async for update in stream: + for content in update.contents: + if content.type == "text": + text_parts.append(content.text) + elif content.type == "function_call": + call_ids.add(content.call_id) + return "".join(text_parts), call_ids, session.session_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["openai", "foundry"]) +async def test_cached_maf_client_isolates_two_concurrent_streaming_runs( + monkeypatch: pytest.MonkeyPatch, + provider: str, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", provider) + if provider == "foundry": + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + transport = _DeterministicResponsesTransport() + chat_client, project_client = _build_chat_client(provider, transport) + manager = MAFClientManager() + builder = "_build_foundry" if provider == "foundry" else "_build_openai" + + try: + with patch.object(MAFClientManager, builder, return_value=chat_client): + shared, _ = manager.build_chat_client_with_target("test-model") + reused, _ = manager.build_chat_client_with_target("test-model") + first_agent = Agent(shared) + second_agent = Agent(reused) + first, second = await asyncio.gather( + _consume_stream(first_agent, "alpha", "session-alpha"), + _consume_stream(second_agent, "beta", "session-beta"), + ) + finally: + await manager.close() + + assert shared is reused is chat_client + assert first == ("reply-alpha", {"call-alpha"}, "session-alpha") + assert second == ("reply-beta", {"call-beta"}, "session-beta") + assert all("beta" not in body for body in transport.requests["alpha"]) + assert all("alpha" not in body for body in transport.requests["beta"]) + assert transport.max_active_requests == 2 + if project_client is not None: + project_client.close.assert_awaited_once_with() + + +async def _run_two_non_streaming_turns( + agent: Agent[Any], + marker: str, + session_id: str, +) -> tuple[str, str, str]: + session = agent.create_session(session_id=session_id) + first = await agent.run(marker, session=session) + second = await agent.run(f"{marker}-followup", session=session) + return first.text, second.text, session.session_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["openai", "foundry"]) +async def test_cached_maf_client_isolates_two_concurrent_non_streaming_runs( + monkeypatch: pytest.MonkeyPatch, + provider: str, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", provider) + if provider == "foundry": + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + transport = _DeterministicResponsesTransport() + chat_client, project_client = _build_chat_client(provider, transport) + manager = MAFClientManager() + builder = "_build_foundry" if provider == "foundry" else "_build_openai" + + try: + with patch.object(MAFClientManager, builder, return_value=chat_client): + shared, _ = manager.build_chat_client_with_target("test-model") + first_agent = Agent(shared) + second_agent = Agent(shared) + first, second = await asyncio.gather( + _run_two_non_streaming_turns(first_agent, "alpha", "session-alpha"), + _run_two_non_streaming_turns(second_agent, "beta", "session-beta"), + ) + finally: + await manager.close() + + assert first == ("reply-alpha", "reply-alpha", "session-alpha") + assert second == ("reply-beta", "reply-beta", "session-beta") + assert len(transport.requests["alpha"]) == 2 + assert len(transport.requests["beta"]) == 2 + assert json.loads(transport.requests["alpha"][1])["previous_response_id"] == "response-alpha" + assert json.loads(transport.requests["beta"][1])["previous_response_id"] == "response-beta" + assert all("beta" not in body for body in transport.requests["alpha"]) + assert all("alpha" not in body for body in transport.requests["beta"]) + assert transport.max_active_requests == 2 + if project_client is not None: + project_client.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["openai", "foundry"]) +async def test_cached_maf_client_isolates_mixed_streaming_and_non_streaming_runs( + monkeypatch: pytest.MonkeyPatch, + provider: str, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", provider) + if provider == "foundry": + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + transport = _DeterministicResponsesTransport() + chat_client, project_client = _build_chat_client(provider, transport) + manager = MAFClientManager() + builder = "_build_foundry" if provider == "foundry" else "_build_openai" + + try: + with patch.object(MAFClientManager, builder, return_value=chat_client): + shared, _ = manager.build_chat_client_with_target("test-model") + first_agent = Agent(shared) + second_agent = Agent(shared) + non_streaming, streaming = await asyncio.gather( + first_agent.run( + "alpha", + session=first_agent.create_session(session_id="session-alpha"), + ), + _consume_stream(second_agent, "beta", "session-beta"), + ) + finally: + await manager.close() + + assert non_streaming.text == "reply-alpha" + assert streaming == ("reply-beta", {"call-beta"}, "session-beta") + assert all("beta" not in body for body in transport.requests["alpha"]) + assert all("alpha" not in body for body in transport.requests["beta"]) + assert transport.max_active_requests == 2 + await manager.close() + if project_client is not None: + project_client.close.assert_awaited_once_with() diff --git a/tests/test_runner_delegation.py b/tests/test_runner_delegation.py index f5cd4fd..b4e195f 100644 --- a/tests/test_runner_delegation.py +++ b/tests/test_runner_delegation.py @@ -32,6 +32,7 @@ from typing import Any, ClassVar import pytest +import pytest_asyncio from agent_framework import MCPStreamableHTTPTool, tool import azure_functions_agents._observability as obs @@ -41,6 +42,7 @@ InferenceTarget, get_client_manager, set_client_manager, + shutdown_client_manager, ) from azure_functions_agents.config.schema import ( BuiltinEndpointsConfig, @@ -158,17 +160,17 @@ def build_chat_client(self, model: str | None) -> Any: return _RunnableFakeChatClient() -@pytest.fixture(autouse=True) -def _restore_client_manager() -> Any: - """Snapshot/restore the process-wide ``ClientManager`` singleton around every test. +@pytest_asyncio.fixture(autouse=True) +async def _restore_client_manager() -> Any: + """Deterministically clean the process-wide manager around every test. ``_build_delegated_agent`` calls ``get_client_manager().build_chat_client(...)``; tests that exercise it install ``_FakeClientManager`` and must not leak that substitution into unrelated tests/modules. """ - original = get_client_manager() + await shutdown_client_manager() yield - set_client_manager(original) + await shutdown_client_manager() class _RecordingSpan: