From da1c78181c0b11aa242d0878845db9b2298f6fd5 Mon Sep 17 00:00:00 2001 From: Tsuyoshi Ushio Date: Wed, 12 Aug 2026 12:06:56 -0700 Subject: [PATCH 1/5] fix: reuse provider clients across agent runs Cache MAF provider clients and credentials per worker process so Functions invocations reuse connection pools and Foundry sessions do not accumulate. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9d876897-6c4c-4e04-ac52-33be4d28676f --- docs/architecture.md | 46 ++- docs/frds/0007-multi-agent-delegation.md | 15 +- src/azure_functions_agents/client_manager.py | 249 ++++++++++-- tests/test_client_manager.py | 379 ++++++++++++++++++- tests/test_runner_delegation.py | 5 +- 5 files changed, 646 insertions(+), 48 deletions(-) diff --git a/docs/architecture.md b/docs/architecture.md index 13c31aea..0f6a1ce6 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 provider clients/connection pools keyed by provider, model, and non-secret endpoint configuration, plus its 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 @@ -347,10 +347,52 @@ To plug in a different chat backend, implement the `ClientManager` interface and 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 serves a request. Replacing +an unused default manager atomically retires it so a thread holding the old +reference cannot create orphaned resources afterward. If the default manager +already owns provider clients, first `await shutdown_client_manager()` and then +call `set_client_manager(...)`; synchronous replacement is rejected rather than +abandoning live connection pools. + 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, resolved model, and relevant +non-secret endpoint/API configuration. 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. + +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/model +configuration is expected to remain stable for the worker lifetime; endpoint +and API-version values participate in the cache key so an embedded host that +changes them does not receive a client for the old endpoint. + +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 the manager currently awaits the pinned implementation's `client.close()` +(`AsyncOpenAI`/httpx) and, for Foundry, the independent +`project_client.close()` (`AIProjectClient`/aiohttp). Missing attributes produce +an explicit warning and tests guard this version-sensitive contract. + +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. + ### 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 86325d71..8743139c 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,7 @@ 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 | ## 6. Test plan @@ -672,6 +674,11 @@ 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). - [ ] 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 254f3350..8bef441e 100644 --- a/src/azure_functions_agents/client_manager.py +++ b/src/azure_functions_agents/client_manager.py @@ -7,27 +7,30 @@ 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 asyncio import os +import threading from abc import ABC, abstractmethod from dataclasses import dataclass +from inspect import isawaitable from typing import Any from ._credential import build_async_credential @@ -47,6 +50,16 @@ class InferenceTarget: model: str | None = None +@dataclass(frozen=True) +class _ClientCacheKey: + provider: str + model: str + endpoint: str = "" + api_version: str = "" + auth_mode: str = "" + organization: str = "" + + class ClientManager(ABC): """Provider-agnostic interface for building chat clients.""" @@ -103,6 +116,13 @@ class MAFClientManager(ClientManager): name = "maf" + def __init__(self) -> None: + self._clients: dict[_ClientCacheKey, Any] = {} + self._async_credential: Any | None = None + self._lock = threading.RLock() + self._closed = False + self._close_task: asyncio.Task[None] | None = None + 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,18 +160,15 @@ def _build_maf_chat_client_with_target( ) -> tuple[Any, InferenceTarget]: provider = self._provider() resolved = self._resolve_model(model, provider) + cache_key = self._client_cache_key(provider, resolved) 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." - ) + with self._lock: + if self._closed: + raise RuntimeError("MAFClientManager is closed and cannot build new clients.") + client = self._clients.get(cache_key) + if client is None: + client = self._build_provider_client(provider, resolved) + self._clients[cache_key] = client return client, InferenceTarget( provider=provider, model=resolved, @@ -161,6 +178,45 @@ def _build_maf_chat_client_with_target( # Internals # ------------------------------------------------------------------ + def _build_provider_client(self, provider: str, model: str) -> Any: + if provider == "openai": + return self._build_openai(model) + if provider == "azure_openai": + return self._build_azure_openai(model) + if provider == "foundry": + return self._build_foundry(model) + raise RuntimeError( + f"Unknown AZURE_FUNCTIONS_AGENTS_PROVIDER '{provider}'. " + "Use one of: openai, azure_openai, foundry." + ) + + @classmethod + def _client_cache_key(cls, provider: str, model: str) -> _ClientCacheKey: + if provider == "openai": + return _ClientCacheKey( + provider, + model, + endpoint=cls._env("OPENAI_BASE_URL"), + auth_mode="api_key", + organization=cls._env("OPENAI_ORG_ID"), + ) + if provider == "azure_openai": + return _ClientCacheKey( + provider, + model, + endpoint=cls._env("AZURE_OPENAI_ENDPOINT"), + api_version=cls._env("AZURE_OPENAI_API_VERSION"), + auth_mode="api_key" if cls._env("AZURE_OPENAI_API_KEY") else "credential", + ) + if provider == "foundry": + return _ClientCacheKey( + provider, + model, + endpoint=cls._env("FOUNDRY_PROJECT_ENDPOINT"), + auth_mode="credential", + ) + return _ClientCacheKey(provider, model) + @staticmethod def _env(name: str) -> str: """Return ``$name`` stripped, or ``""`` if missing/blank. @@ -200,11 +256,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) -> Any: 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 +272,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) -> Any: 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,8 +294,112 @@ 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) -> Any: + with self._lock: + if self._closed: + raise RuntimeError("MAFClientManager is closed and cannot build a credential.") + if self._async_credential is None: + self._async_credential = build_async_credential() + return self._async_credential + + async def close(self) -> None: + """Close all provider transports and the shared credential exactly once.""" + with self._lock: + close_task = self._close_task + if close_task is None: + self._closed = True + cached_clients = list(self._clients.items()) + self._clients.clear() + credential = self._async_credential + self._async_credential = None + close_task = asyncio.create_task( + self._close_detached_resources(cached_clients, credential) + ) + self._close_task = close_task + await asyncio.shield(close_task) + + async def _close_detached_resources( + self, + cached_clients: list[tuple[_ClientCacheKey, Any]], + credential: Any, + ) -> None: + errors: list[Exception] = [] + closed_resource_ids: set[int] = set() + for key, chat_client in cached_clients: + await self._close_owned_resource( + getattr(chat_client, "client", None), + f"{key.provider} AsyncOpenAI transport", + errors, + closed_resource_ids, + ) + if key.provider == "foundry": + await self._close_owned_resource( + getattr(chat_client, "project_client", None), + "Foundry AIProjectClient transport", + errors, + closed_resource_ids, + ) + await self._close_owned_resource( + credential, + "shared async credential", + errors, + closed_resource_ids, + required=False, ) + if errors: + raise ExceptionGroup("Failed to close one or more MAF client resources.", errors) + + def _has_owned_resources(self) -> bool: + with self._lock: + return bool( + self._clients + or self._async_credential is not None + or (self._close_task is not None and not self._close_task.done()) + ) + + def _retire_if_unused(self) -> bool: + """Atomically prevent future builds when no owned resource exists.""" + with self._lock: + if self._has_owned_resources(): + return False + self._closed = True + return True + + @staticmethod + async def _close_owned_resource( + resource: Any, + label: str, + errors: list[Exception], + closed_resource_ids: set[int], + *, + required: bool = True, + ) -> None: + if resource is None: + if required: + logger.warning( + "%s is unavailable; agent-framework 1.3 client internals may have changed.", + label, + ) + return + resource_id = id(resource) + if resource_id in closed_resource_ids: + return + closed_resource_ids.add(resource_id) + close = getattr(resource, "close", None) + if not callable(close): + if required: + logger.warning("%s does not expose close().", label) + return + try: + result = close() + if isawaitable(result): + await result + except Exception as exc: + logger.error("Failed to close %s: %s", label, exc) + errors.append(exc) # --------------------------------------------------------------------------- @@ -249,6 +407,7 @@ def _build_foundry(cls, model: str) -> Any: # --------------------------------------------------------------------------- _INSTANCE: ClientManager | None = None +_INSTANCE_LOCK = threading.Lock() def get_client_manager() -> ClientManager: @@ -259,27 +418,39 @@ 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 None: + _INSTANCE = MAFClientManager() + logger.info("ClientManager initialized: %s", _INSTANCE.name) + 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. Replacing the default manager retires it atomically when unused. + Once it owns provider resources, callers must await + :func:`shutdown_client_manager` before installing a replacement. """ global _INSTANCE - _INSTANCE = manager + with _INSTANCE_LOCK: + current = _INSTANCE + if current is manager: + return + if isinstance(current, MAFClientManager) and not current._retire_if_unused(): + raise RuntimeError( + "The active MAFClientManager owns provider resources. " + "Await shutdown_client_manager() before replacing it." + ) + _INSTANCE = manager async def shutdown_client_manager() -> None: """Close the active manager (if any). Idempotent.""" global _INSTANCE - if _INSTANCE is not None: - try: - await _INSTANCE.close() - finally: - _INSTANCE = None + with _INSTANCE_LOCK: + manager = _INSTANCE + _INSTANCE = None + if manager is not None: + await manager.close() diff --git a/tests/test_client_manager.py b/tests/test_client_manager.py index 22d3f7e3..7acd14a2 100644 --- a/tests/test_client_manager.py +++ b/tests/test_client_manager.py @@ -1,7 +1,11 @@ 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 @@ -12,6 +16,9 @@ ClientManager, InferenceTarget, MAFClientManager, + get_client_manager, + set_client_manager, + shutdown_client_manager, ) @@ -158,6 +165,376 @@ 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 = object() + + 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_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 = [object(), object()] + + 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 + + +def test_maf_manager_partitions_cached_clients_by_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project-one.example") + manager = MAFClientManager() + clients = [object(), object()] + + with patch.object(MAFClientManager, "_build_foundry", side_effect=clients) as build: + first, _ = manager.build_chat_client_with_target("shared-model") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project-two.example") + second, _ = manager.build_chat_client_with_target("shared-model") + + assert first is clients[0] + assert second is clients[1] + assert build.call_count == 2 + + +def test_maf_manager_partitions_azure_openai_clients_by_api_version( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "azure_openai") + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://account.openai.azure.com") + monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key") + monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "version-one") + manager = MAFClientManager() + clients = [object(), object()] + + with patch.object(MAFClientManager, "_build_azure_openai", side_effect=clients) as build: + first, _ = manager.build_chat_client_with_target("shared-model") + monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "version-two") + second, _ = manager.build_chat_client_with_target("shared-model") + + assert first is clients[0] + assert second is clients[1] + assert build.call_count == 2 + + +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 = object() + + 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_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") + + +@pytest.mark.asyncio +async def test_maf_manager_cleanup_survives_caller_cancellation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + manager = MAFClientManager() + close_started = asyncio.Event() + allow_close = asyncio.Event() + + async def _slow_close() -> None: + close_started.set() + await allow_close.wait() + + openai_client = SimpleNamespace(close=AsyncMock(side_effect=_slow_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") + first_close = asyncio.create_task(manager.close()) + await close_started.wait() + first_close.cancel() + with pytest.raises(asyncio.CancelledError): + await first_close + + allow_close.set() + 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_shutdown_does_not_clear_manager_installed_while_old_manager_closes() -> None: + get_client_manager() + 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() + + set_client_manager(replacement) + allow_close.set() + await shutdown + + assert get_client_manager() is replacement + set_client_manager(MAFClientManager()) + + +@pytest.mark.asyncio +async def test_set_client_manager_rejects_abandoning_owned_provider_resources( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") + get_client_manager() + manager = MAFClientManager() + chat_client = SimpleNamespace(client=SimpleNamespace(close=AsyncMock())) + + with patch.object(MAFClientManager, "_build_openai", return_value=chat_client): + manager.build_chat_client_with_target("shared-model") + set_client_manager(manager) + + with pytest.raises(RuntimeError, match="owns provider resources"): + set_client_manager(MAFClientManager()) + + await shutdown_client_manager() + set_client_manager(MAFClientManager()) + + +def test_set_client_manager_retires_displaced_unused_maf_manager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") + displaced = MAFClientManager() + replacement = MAFClientManager() + set_client_manager(displaced) + + set_client_manager(replacement) + + with ( + patch.object(MAFClientManager, "_build_openai", return_value=object()), + pytest.raises(RuntimeError, match="closed"), + ): + displaced.build_chat_client_with_target("model-one") + + def test_custom_manager_target_fallback_builds_client_once() -> None: class CustomManager(ClientManager): calls = 0 diff --git a/tests/test_runner_delegation.py b/tests/test_runner_delegation.py index f5cd4fd6..a2162040 100644 --- a/tests/test_runner_delegation.py +++ b/tests/test_runner_delegation.py @@ -39,6 +39,7 @@ from azure_functions_agents.client_manager import ( ClientManager, InferenceTarget, + MAFClientManager, get_client_manager, set_client_manager, ) @@ -166,9 +167,9 @@ def _restore_client_manager() -> Any: tests that exercise it install ``_FakeClientManager`` and must not leak that substitution into unrelated tests/modules. """ - original = get_client_manager() + get_client_manager() yield - set_client_manager(original) + set_client_manager(MAFClientManager()) class _RecordingSpan: From 35c7297f0bd8c1c52240d7834689e01f8880e525 Mon Sep 17 00:00:00 2001 From: Tsuyoshi Ushio Date: Wed, 12 Aug 2026 15:22:34 -0700 Subject: [PATCH 2/5] fix: simplify provider client lifecycle Use direct awaited cleanup and one shutdown sentinel while preserving process-wide provider connection reuse. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9d876897-6c4c-4e04-ac52-33be4d28676f --- README.md | 41 ++++ docs/architecture.md | 27 +-- docs/frds/0007-multi-agent-delegation.md | 1 + src/azure_functions_agents/client_manager.py | 110 +++++----- tests/test_client_manager.py | 207 ++++++++++++------- tests/test_runner_delegation.py | 13 +- 6 files changed, 245 insertions(+), 154 deletions(-) diff --git a/README.md b/README.md index 69ef8422..9a3acbce 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,47 @@ 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. + +Keep provider endpoint, API-version, organization, and authentication settings +stable for a worker lifetime. If one changes, the runtime rejects the next +client request; restart the Functions worker to apply the new configuration. + +### Plugging in a custom client manager + +Implement `ClientManager` and install it once, before the default manager is +first requested: + +```python +from azure_functions_agents import ClientManager, set_client_manager + + +class MyClientManager(ClientManager): + def resolve_model(self, requested: str | None) -> str: + return requested or "my-default-model" + + def build_chat_client(self, model: str | None): + return build_my_chat_client(self.resolve_model(model)) + + +set_client_manager(MyClientManager()) +``` + +The active manager cannot be replaced synchronously, whether it is the default +or a custom implementation. Tests and embedding hosts that own an async +lifecycle must first `await shutdown_client_manager()`, then call +`set_client_manager(...)`. Azure Functions does not currently expose a supported +async worker-shutdown hook, so the runtime keeps its default clients for the +worker lifetime rather than attempting cleanup from `atexit` or a signal handler; +the host gap is tracked in +[azure-functions-python-worker#1904](https://github.com/Azure/azure-functions-python-worker/issues/1904). + ## Quick Start ### 1. Create the agent file diff --git a/docs/architecture.md b/docs/architecture.md index 0f6a1ce6..97115aab 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. The default manager owns a worker-process cache of provider clients/connection pools keyed by provider, model, and non-secret endpoint configuration, plus its shared async credential. | `ClientManager`, `InferenceTarget`, `get_client_manager()`, `set_client_manager()`, `shutdown_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 provider clients/connection pools keyed by provider and model, enforces stable provider configuration for that worker lifetime, and owns 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` | @@ -347,12 +347,11 @@ To plug in a different chat backend, implement the `ClientManager` interface and 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 serves a request. Replacing -an unused default manager atomically retires it so a thread holding the old -reference cannot create orphaned resources afterward. If the default manager -already owns provider clients, first `await shutdown_client_manager()` and then -call `set_client_manager(...)`; synchronous replacement is rejected rather than -abandoning live connection pools. +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. @@ -361,8 +360,7 @@ The runner calls `build_chat_client_with_target()` and receives the client plus ### 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, resolved model, and relevant -non-secret endpoint/API configuration. Primary agents, declared triggers, +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 @@ -372,10 +370,11 @@ 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/model -configuration is expected to remain stable for the worker lifetime; endpoint -and API-version values participate in the cache key so an embedded host that -changes them does not receive a client for the old endpoint. +lifetime boundary for the cached aiohttp/httpx transports. Provider +configuration is immutable after that provider's first client is built: +endpoint, API version, organization, and authentication-mode changes are +rejected with an instruction to restart the worker. Different resolved models +may still share the same stable provider configuration. Managed-identity-backed Azure OpenAI and Foundry clients share one manager-owned async credential and token cache. `MAFClientManager.close()` is @@ -392,6 +391,8 @@ 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 diff --git a/docs/frds/0007-multi-agent-delegation.md b/docs/frds/0007-multi-agent-delegation.md index 8743139c..e2e0157a 100644 --- a/docs/frds/0007-multi-agent-delegation.md +++ b/docs/frds/0007-multi-agent-delegation.md @@ -641,6 +641,7 @@ handoff participant may itself declare `subagents` and delegate. | 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 | ## 6. Test plan diff --git a/src/azure_functions_agents/client_manager.py b/src/azure_functions_agents/client_manager.py index 8bef441e..323fc5ca 100644 --- a/src/azure_functions_agents/client_manager.py +++ b/src/azure_functions_agents/client_manager.py @@ -25,7 +25,6 @@ from __future__ import annotations -import asyncio import os import threading from abc import ABC, abstractmethod @@ -54,6 +53,10 @@ class InferenceTarget: class _ClientCacheKey: provider: str model: str + + +@dataclass(frozen=True) +class _ProviderConfig: endpoint: str = "" api_version: str = "" auth_mode: str = "" @@ -118,10 +121,10 @@ class MAFClientManager(ClientManager): def __init__(self) -> None: self._clients: dict[_ClientCacheKey, Any] = {} + self._provider_configs: dict[str, _ProviderConfig] = {} self._async_credential: Any | None = None self._lock = threading.RLock() self._closed = False - self._close_task: asyncio.Task[None] | None = None def resolve_model(self, requested: str | None) -> str: """Resolve model as requested > provider-specific env > runtime env > default.""" @@ -160,15 +163,25 @@ def _build_maf_chat_client_with_target( ) -> tuple[Any, InferenceTarget]: provider = self._provider() resolved = self._resolve_model(model, provider) - cache_key = self._client_cache_key(provider, resolved) - logger.info("MAF provider=%s model=%s", provider, resolved) + cache_key = _ClientCacheKey(provider, resolved) with self._lock: if self._closed: raise RuntimeError("MAFClientManager is closed and cannot build new clients.") + provider_config = self._provider_config(provider) + previous_config = self._provider_configs.get(provider) + if previous_config is not None and previous_config != provider_config: + raise RuntimeError( + f"{provider} provider configuration changed during this worker lifetime; " + "restart the worker to apply updated settings" + ) client = self._clients.get(cache_key) if client is None: client = self._build_provider_client(provider, resolved) + self._provider_configs[provider] = provider_config self._clients[cache_key] = client + 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 client, InferenceTarget( provider=provider, model=resolved, @@ -191,31 +204,25 @@ def _build_provider_client(self, provider: str, model: str) -> Any: ) @classmethod - def _client_cache_key(cls, provider: str, model: str) -> _ClientCacheKey: + def _provider_config(cls, provider: str) -> _ProviderConfig: if provider == "openai": - return _ClientCacheKey( - provider, - model, + return _ProviderConfig( endpoint=cls._env("OPENAI_BASE_URL"), auth_mode="api_key", organization=cls._env("OPENAI_ORG_ID"), ) if provider == "azure_openai": - return _ClientCacheKey( - provider, - model, + return _ProviderConfig( endpoint=cls._env("AZURE_OPENAI_ENDPOINT"), api_version=cls._env("AZURE_OPENAI_API_VERSION"), auth_mode="api_key" if cls._env("AZURE_OPENAI_API_KEY") else "credential", ) if provider == "foundry": - return _ClientCacheKey( - provider, - model, + return _ProviderConfig( endpoint=cls._env("FOUNDRY_PROJECT_ENDPOINT"), auth_mode="credential", ) - return _ClientCacheKey(provider, model) + return _ProviderConfig() @staticmethod def _env(name: str) -> str: @@ -308,24 +315,15 @@ def _get_async_credential(self) -> Any: async def close(self) -> None: """Close all provider transports and the shared credential exactly once.""" with self._lock: - close_task = self._close_task - if close_task is None: - self._closed = True - cached_clients = list(self._clients.items()) - self._clients.clear() - credential = self._async_credential - self._async_credential = None - close_task = asyncio.create_task( - self._close_detached_resources(cached_clients, credential) - ) - self._close_task = close_task - await asyncio.shield(close_task) + if self._closed: + return + self._closed = True + cached_clients = list(self._clients.items()) + self._clients.clear() + self._provider_configs.clear() + credential = self._async_credential + self._async_credential = None - async def _close_detached_resources( - self, - cached_clients: list[tuple[_ClientCacheKey, Any]], - credential: Any, - ) -> None: errors: list[Exception] = [] closed_resource_ids: set[int] = set() for key, chat_client in cached_clients: @@ -352,22 +350,6 @@ async def _close_detached_resources( if errors: raise ExceptionGroup("Failed to close one or more MAF client resources.", errors) - def _has_owned_resources(self) -> bool: - with self._lock: - return bool( - self._clients - or self._async_credential is not None - or (self._close_task is not None and not self._close_task.done()) - ) - - def _retire_if_unused(self) -> bool: - """Atomically prevent future builds when no owned resource exists.""" - with self._lock: - if self._has_owned_resources(): - return False - self._closed = True - return True - @staticmethod async def _close_owned_resource( resource: Any, @@ -406,7 +388,8 @@ async def _close_owned_resource( # Process-wide singleton selection # --------------------------------------------------------------------------- -_INSTANCE: ClientManager | None = None +_SHUTTING_DOWN = object() +_INSTANCE: ClientManager | object | None = None _INSTANCE_LOCK = threading.Lock() @@ -419,9 +402,12 @@ def get_client_manager() -> ClientManager: """ global _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 @@ -429,28 +415,38 @@ 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. Replacing the default manager retires it atomically when unused. - Once it owns provider resources, callers must await - :func:`shutdown_client_manager` before installing a replacement. + backend. Call and await :func:`shutdown_client_manager` before replacing + any active manager, including a custom implementation. """ global _INSTANCE with _INSTANCE_LOCK: + if _INSTANCE is _SHUTTING_DOWN: + raise RuntimeError("Client manager shutdown is in progress") current = _INSTANCE if current is manager: return - if isinstance(current, MAFClientManager) and not current._retire_if_unused(): + if current is not None: raise RuntimeError( - "The active MAFClientManager owns provider resources. " + "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 with _INSTANCE_LOCK: manager = _INSTANCE - _INSTANCE = None - if manager is not None: + 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 7acd14a2..54fe11d1 100644 --- a/tests/test_client_manager.py +++ b/tests/test_client_manager.py @@ -8,6 +8,7 @@ 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 ( @@ -22,6 +23,13 @@ ) +@pytest_asyncio.fixture(autouse=True) +async def _reset_process_client_manager() -> None: + await shutdown_client_manager() + yield + await shutdown_client_manager() + + @pytest.mark.parametrize( ("provider", "provider_env", "provider_model"), [ @@ -185,6 +193,68 @@ async def test_maf_manager_reuses_provider_client_on_one_worker_loop( 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=object()): + 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 = object() + + 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: @@ -203,25 +273,25 @@ def test_maf_manager_partitions_cached_clients_by_resolved_model( assert build.call_count == 2 -def test_maf_manager_partitions_cached_clients_by_endpoint( +def test_maf_manager_rejects_endpoint_change_during_worker_lifetime( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project-one.example") manager = MAFClientManager() - clients = [object(), object()] + client = object() - with patch.object(MAFClientManager, "_build_foundry", side_effect=clients) as build: + with patch.object(MAFClientManager, "_build_foundry", return_value=client) as build: first, _ = manager.build_chat_client_with_target("shared-model") monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project-two.example") - second, _ = manager.build_chat_client_with_target("shared-model") + with pytest.raises(RuntimeError, match="provider configuration changed"): + manager.build_chat_client_with_target("shared-model") - assert first is clients[0] - assert second is clients[1] - assert build.call_count == 2 + assert first is client + build.assert_called_once_with("shared-model") -def test_maf_manager_partitions_azure_openai_clients_by_api_version( +def test_maf_manager_rejects_api_version_change_during_worker_lifetime( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "azure_openai") @@ -229,16 +299,16 @@ def test_maf_manager_partitions_azure_openai_clients_by_api_version( monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key") monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "version-one") manager = MAFClientManager() - clients = [object(), object()] + client = object() - with patch.object(MAFClientManager, "_build_azure_openai", side_effect=clients) as build: + with patch.object(MAFClientManager, "_build_azure_openai", return_value=client) as build: first, _ = manager.build_chat_client_with_target("shared-model") monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "version-two") - second, _ = manager.build_chat_client_with_target("shared-model") + with pytest.raises(RuntimeError, match="provider configuration changed"): + manager.build_chat_client_with_target("shared-model") - assert first is clients[0] - assert second is clients[1] - assert build.call_count == 2 + assert first is client + build.assert_called_once_with("shared-model") def test_maf_manager_publishes_one_client_during_concurrent_first_use( @@ -426,50 +496,23 @@ async def test_maf_manager_rejects_build_after_close( manager.build_chat_client_with_target("model-one") -@pytest.mark.asyncio -async def test_maf_manager_cleanup_survives_caller_cancellation( +def test_maf_manager_close_runs_from_sync_embedding_boundary( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") - monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") manager = MAFClientManager() - close_started = asyncio.Event() - allow_close = asyncio.Event() - - async def _slow_close() -> None: - close_started.set() - await allow_close.wait() - - openai_client = SimpleNamespace(close=AsyncMock(side_effect=_slow_close)) - project_client = SimpleNamespace(close=AsyncMock()) - chat_client = SimpleNamespace(client=openai_client, project_client=project_client) - credential = SimpleNamespace(close=AsyncMock()) + transport = SimpleNamespace(close=AsyncMock()) + chat_client = SimpleNamespace(client=transport) - with ( - patch("agent_framework.foundry.FoundryChatClient", return_value=chat_client), - patch( - "azure_functions_agents.client_manager.build_async_credential", - return_value=credential, - ), - ): + with patch.object(MAFClientManager, "_build_openai", return_value=chat_client): manager.build_chat_client_with_target("shared-model") - first_close = asyncio.create_task(manager.close()) - await close_started.wait() - first_close.cancel() - with pytest.raises(asyncio.CancelledError): - await first_close - - allow_close.set() - await manager.close() + asyncio.run(manager.close()) - openai_client.close.assert_awaited_once_with() - project_client.close.assert_awaited_once_with() - credential.close.assert_awaited_once_with() + transport.close.assert_awaited_once_with() @pytest.mark.asyncio -async def test_shutdown_does_not_clear_manager_installed_while_old_manager_closes() -> None: - get_client_manager() +async def test_shutdown_blocks_get_and_set_until_cleanup_finishes() -> None: close_started = asyncio.Event() allow_close = asyncio.Event() @@ -490,49 +533,57 @@ async def close(self) -> None: shutdown = asyncio.create_task(shutdown_client_manager()) await close_started.wait() - set_client_manager(replacement) + 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 - set_client_manager(MAFClientManager()) @pytest.mark.asyncio -async def test_set_client_manager_rejects_abandoning_owned_provider_resources( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") - get_client_manager() - manager = MAFClientManager() - chat_client = SimpleNamespace(client=SimpleNamespace(close=AsyncMock())) +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" - with patch.object(MAFClientManager, "_build_openai", return_value=chat_client): - manager.build_chat_client_with_target("shared-model") - set_client_manager(manager) + def build_chat_client(self, model: str | None) -> Any: + return object() - with pytest.raises(RuntimeError, match="owns provider resources"): - set_client_manager(MAFClientManager()) + async def close(self) -> None: + raise RuntimeError("close failed") - await shutdown_client_manager() - set_client_manager(MAFClientManager()) + set_client_manager(FailingManager()) + with pytest.raises(RuntimeError, match="close failed"): + await shutdown_client_manager() -def test_set_client_manager_retires_displaced_unused_maf_manager( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "openai") - displaced = MAFClientManager() - replacement = MAFClientManager() - set_client_manager(displaced) + assert isinstance(get_client_manager(), MAFClientManager) - set_client_manager(replacement) - with ( - patch.object(MAFClientManager, "_build_openai", return_value=object()), - pytest.raises(RuntimeError, match="closed"), - ): - displaced.build_chat_client_with_target("model-one") +@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: diff --git a/tests/test_runner_delegation.py b/tests/test_runner_delegation.py index a2162040..b4e195fe 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 @@ -39,9 +40,9 @@ from azure_functions_agents.client_manager import ( ClientManager, InferenceTarget, - MAFClientManager, get_client_manager, set_client_manager, + shutdown_client_manager, ) from azure_functions_agents.config.schema import ( BuiltinEndpointsConfig, @@ -159,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. """ - get_client_manager() + await shutdown_client_manager() yield - set_client_manager(MAFClientManager()) + await shutdown_client_manager() class _RecordingSpan: From 4613ab1eba5fc29f4fcc4abeef80d90168668a05 Mon Sep 17 00:00:00 2001 From: Tsuyoshi Ushio Date: Wed, 12 Aug 2026 20:14:45 -0700 Subject: [PATCH 3/5] fix: type provider client ownership Record concrete MAF transports in typed cache entries and treat provider settings as process configuration. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9d876897-6c4c-4e04-ac52-33be4d28676f --- README.md | 39 ++-- docs/architecture.md | 21 +- docs/frds/0007-multi-agent-delegation.md | 1 + src/azure_functions_agents/client_manager.py | 195 ++++++++----------- tests/test_client_manager.py | 115 +++++++---- 5 files changed, 179 insertions(+), 192 deletions(-) diff --git a/README.md b/README.md index 9a3acbce..d6e70929 100644 --- a/README.md +++ b/README.md @@ -50,37 +50,20 @@ 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. -Keep provider endpoint, API-version, organization, and authentication settings -stable for a worker lifetime. If one changes, the runtime rejects the next -client request; restart the Functions worker to apply the new configuration. +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. -### Plugging in a custom client manager +### Advanced client manager lifecycle -Implement `ClientManager` and install it once, before the default manager is -first requested: +`ClientManager` is an advanced extension point used by test fakes, custom +MAF-compatible provider clients, and embedding hosts. Install a custom manager +once, before the default manager is requested. Replacing any active manager +requires `await shutdown_client_manager()` before `set_client_manager(...)`. -```python -from azure_functions_agents import ClientManager, set_client_manager - - -class MyClientManager(ClientManager): - def resolve_model(self, requested: str | None) -> str: - return requested or "my-default-model" - - def build_chat_client(self, model: str | None): - return build_my_chat_client(self.resolve_model(model)) - - -set_client_manager(MyClientManager()) -``` - -The active manager cannot be replaced synchronously, whether it is the default -or a custom implementation. Tests and embedding hosts that own an async -lifecycle must first `await shutdown_client_manager()`, then call -`set_client_manager(...)`. Azure Functions does not currently expose a supported -async worker-shutdown hook, so the runtime keeps its default clients for the -worker lifetime rather than attempting cleanup from `atexit` or a signal handler; -the host gap is tracked in +Azure Functions does not currently expose a supported async worker-shutdown +hook, so the runtime keeps its default clients for the worker lifetime rather +than attempting cleanup from `atexit` or a signal handler. The host gap is tracked in [azure-functions-python-worker#1904](https://github.com/Azure/azure-functions-python-worker/issues/1904). ## Quick Start diff --git a/docs/architecture.md b/docs/architecture.md index 97115aab..a78009e1 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. The default manager owns a worker-process cache of provider clients/connection pools keyed by provider and model, enforces stable provider configuration for that worker lifetime, and owns a shared async credential. | `ClientManager`, `InferenceTarget`, `get_client_manager()`, `set_client_manager()`, `shutdown_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` | @@ -343,7 +343,7 @@ 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](../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 the README section [Advanced client manager lifecycle](../README.md#advanced-client-manager-lifecycle). 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. @@ -370,20 +370,19 @@ 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 -configuration is immutable after that provider's first client is built: -endpoint, API version, organization, and authentication-mode changes are -rejected with an instruction to restart the worker. Different resolved models -may still share the same stable provider configuration. +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 the manager currently awaits the pinned implementation's `client.close()` -(`AsyncOpenAI`/httpx) and, for Foundry, the independent -`project_client.close()` (`AIProjectClient`/aiohttp). Missing attributes produce -an explicit warning and tests guard this version-sensitive contract. +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 diff --git a/docs/frds/0007-multi-agent-delegation.md b/docs/frds/0007-multi-agent-delegation.md index e2e0157a..0f5ecb5f 100644 --- a/docs/frds/0007-multi-agent-delegation.md +++ b/docs/frds/0007-multi-agent-delegation.md @@ -642,6 +642,7 @@ handoff participant may itself declare `subagents` and delegate. | 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 | ## 6. Test plan diff --git a/src/azure_functions_agents/client_manager.py b/src/azure_functions_agents/client_manager.py index 323fc5ca..357ffa48 100644 --- a/src/azure_functions_agents/client_manager.py +++ b/src/azure_functions_agents/client_manager.py @@ -29,13 +29,18 @@ import threading from abc import ABC, abstractmethod from dataclasses import dataclass -from inspect import isawaitable -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 # --------------------------------------------------------------------------- @@ -55,12 +60,36 @@ class _ClientCacheKey: model: str +class _AsyncCloseable(Protocol): + async def close(self) -> None: ... + + +@dataclass(frozen=True) +class _OwnedResource: + value: _AsyncCloseable + label: str + + @dataclass(frozen=True) -class _ProviderConfig: - endpoint: str = "" - api_version: str = "" - auth_mode: str = "" - organization: str = "" +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): @@ -120,9 +149,8 @@ class MAFClientManager(ClientManager): name = "maf" def __init__(self) -> None: - self._clients: dict[_ClientCacheKey, Any] = {} - self._provider_configs: dict[str, _ProviderConfig] = {} - self._async_credential: Any | None = None + self._clients: dict[_ClientCacheKey, _ManagedChatClient] = {} + self._async_credential: AsyncDefaultAzureCredential | None = None self._lock = threading.RLock() self._closed = False @@ -165,24 +193,15 @@ def _build_maf_chat_client_with_target( resolved = self._resolve_model(model, provider) cache_key = _ClientCacheKey(provider, resolved) with self._lock: - if self._closed: - raise RuntimeError("MAFClientManager is closed and cannot build new clients.") - provider_config = self._provider_config(provider) - previous_config = self._provider_configs.get(provider) - if previous_config is not None and previous_config != provider_config: - raise RuntimeError( - f"{provider} provider configuration changed during this worker lifetime; " - "restart the worker to apply updated settings" - ) - client = self._clients.get(cache_key) - if client is None: - client = self._build_provider_client(provider, resolved) - self._provider_configs[provider] = provider_config - self._clients[cache_key] = client + 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 client, InferenceTarget( + return managed.client, InferenceTarget( provider=provider, model=resolved, ) @@ -191,39 +210,36 @@ def _build_maf_chat_client_with_target( # Internals # ------------------------------------------------------------------ - def _build_provider_client(self, provider: str, model: str) -> Any: + def _build_provider_client(self, provider: str, model: str) -> _ManagedChatClient: if provider == "openai": - return self._build_openai(model) + client = self._build_openai(model) + return _ManagedChatClient( + client, + (_OwnedResource(client.client, "OpenAI AsyncOpenAI transport"),), + ) if provider == "azure_openai": - return self._build_azure_openai(model) + client = self._build_azure_openai(model) + return _ManagedChatClient( + client, + (_OwnedResource(client.client, "Azure OpenAI AsyncOpenAI transport"),), + ) if provider == "foundry": - return self._build_foundry(model) + 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." ) - @classmethod - def _provider_config(cls, provider: str) -> _ProviderConfig: - if provider == "openai": - return _ProviderConfig( - endpoint=cls._env("OPENAI_BASE_URL"), - auth_mode="api_key", - organization=cls._env("OPENAI_ORG_ID"), - ) - if provider == "azure_openai": - return _ProviderConfig( - endpoint=cls._env("AZURE_OPENAI_ENDPOINT"), - api_version=cls._env("AZURE_OPENAI_API_VERSION"), - auth_mode="api_key" if cls._env("AZURE_OPENAI_API_KEY") else "credential", - ) - if provider == "foundry": - return _ProviderConfig( - endpoint=cls._env("FOUNDRY_PROJECT_ENDPOINT"), - auth_mode="credential", - ) - return _ProviderConfig() - @staticmethod def _env(name: str) -> str: """Return ``$name`` stripped, or ``""`` if missing/blank. @@ -255,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( @@ -263,7 +279,7 @@ def _build_openai(cls, model: str) -> Any: api_key=cls._env("OPENAI_API_KEY") or None, ) - def _build_azure_openai(self, model: str) -> Any: + def _build_azure_openai(self, model: str) -> OpenAIChatClient: from agent_framework.openai import OpenAIChatClient endpoint = self._env("AZURE_OPENAI_ENDPOINT") @@ -289,7 +305,7 @@ def _build_azure_openai(self, model: str) -> Any: kwargs["credential"] = self._get_async_credential() return OpenAIChatClient(**kwargs) - def _build_foundry(self, model: str) -> Any: + def _build_foundry(self, model: str) -> FoundryChatClient: from agent_framework.foundry import FoundryChatClient endpoint = self._env("FOUNDRY_PROJECT_ENDPOINT") @@ -304,85 +320,40 @@ def _build_foundry(self, model: str) -> Any: credential=self._get_async_credential(), ) - def _get_async_credential(self) -> Any: + def _get_async_credential(self) -> AsyncDefaultAzureCredential: with self._lock: - if self._closed: - raise RuntimeError("MAFClientManager is closed and cannot build a credential.") + 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.items()) + cached_clients = list(self._clients.values()) self._clients.clear() - self._provider_configs.clear() credential = self._async_credential self._async_credential = None errors: list[Exception] = [] - closed_resource_ids: set[int] = set() - for key, chat_client in cached_clients: - await self._close_owned_resource( - getattr(chat_client, "client", None), - f"{key.provider} AsyncOpenAI transport", - errors, - closed_resource_ids, - ) - if key.provider == "foundry": - await self._close_owned_resource( - getattr(chat_client, "project_client", None), - "Foundry AIProjectClient transport", - errors, - closed_resource_ids, + 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"),) ) - await self._close_owned_resource( - credential, - "shared async credential", - errors, - closed_resource_ids, - required=False, - ) + ) if errors: raise ExceptionGroup("Failed to close one or more MAF client resources.", errors) - @staticmethod - async def _close_owned_resource( - resource: Any, - label: str, - errors: list[Exception], - closed_resource_ids: set[int], - *, - required: bool = True, - ) -> None: - if resource is None: - if required: - logger.warning( - "%s is unavailable; agent-framework 1.3 client internals may have changed.", - label, - ) - return - resource_id = id(resource) - if resource_id in closed_resource_ids: - return - closed_resource_ids.add(resource_id) - close = getattr(resource, "close", None) - if not callable(close): - if required: - logger.warning("%s does not expose close().", label) - return - try: - result = close() - if isawaitable(result): - await result - except Exception as exc: - logger.error("Failed to close %s: %s", label, exc) - errors.append(exc) - # --------------------------------------------------------------------------- # Process-wide singleton selection diff --git a/tests/test_client_manager.py b/tests/test_client_manager.py index 54fe11d1..3cbc68d9 100644 --- a/tests/test_client_manager.py +++ b/tests/test_client_manager.py @@ -30,6 +30,16 @@ async def _reset_process_client_manager() -> None: 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"), [ @@ -140,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") @@ -155,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, @@ -180,7 +190,7 @@ async def test_maf_manager_reuses_provider_client_on_one_worker_loop( monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") manager = MAFClientManager() - client = object() + 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") @@ -201,7 +211,11 @@ def test_maf_manager_logs_cache_creation_at_info_and_hit_at_debug( manager = MAFClientManager() caplog.set_level("DEBUG", logger="azure.functions.AgentRuntime") - with patch.object(MAFClientManager, "_build_openai", return_value=object()): + 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") @@ -243,7 +257,7 @@ def test_maf_manager_reuses_auto_detected_provider_client( monkeypatch.delenv("OPENAI_API_KEY", raising=False) monkeypatch.setenv(endpoint_name, endpoint) manager = MAFClientManager() - client = object() + 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") @@ -261,7 +275,10 @@ def test_maf_manager_partitions_cached_clients_by_resolved_model( monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") manager = MAFClientManager() - clients = [object(), object()] + 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") @@ -273,42 +290,32 @@ def test_maf_manager_partitions_cached_clients_by_resolved_model( assert build.call_count == 2 -def test_maf_manager_rejects_endpoint_change_during_worker_lifetime( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") - monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project-one.example") - manager = MAFClientManager() - client = object() - - with patch.object(MAFClientManager, "_build_foundry", return_value=client) as build: - first, _ = manager.build_chat_client_with_target("shared-model") - monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project-two.example") - with pytest.raises(RuntimeError, match="provider configuration changed"): - manager.build_chat_client_with_target("shared-model") - - assert first is client - build.assert_called_once_with("shared-model") - - -def test_maf_manager_rejects_api_version_change_during_worker_lifetime( +@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", "azure_openai") - monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://account.openai.azure.com") - monkeypatch.setenv("AZURE_OPENAI_API_KEY", "test-key") - monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "version-one") - manager = MAFClientManager() - client = object() - - with patch.object(MAFClientManager, "_build_azure_openai", return_value=client) as build: - first, _ = manager.build_chat_client_with_target("shared-model") - monkeypatch.setenv("AZURE_OPENAI_API_VERSION", "version-two") - with pytest.raises(RuntimeError, match="provider configuration changed"): - manager.build_chat_client_with_target("shared-model") + monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", provider) + if provider == "foundry": + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") - assert first is client - build.assert_called_once_with("shared-model") + 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( @@ -317,7 +324,7 @@ def test_maf_manager_publishes_one_client_during_concurrent_first_use( monkeypatch.setenv("AZURE_FUNCTIONS_AGENTS_PROVIDER", "foundry") monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://project.example") manager = MAFClientManager() - client = object() + client = _fake_provider_client("foundry") def _build(_model: str) -> object: time.sleep(0.02) @@ -385,6 +392,32 @@ async def test_maf_manager_closes_openai_transport_once( 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, @@ -610,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 From cf11cbf89e99f66a6e6381cda07ee3243e3989df Mon Sep 17 00:00:00 2001 From: Tsuyoshi Ushio Date: Thu, 13 Aug 2026 10:15:00 -0700 Subject: [PATCH 4/5] test: validate shared client concurrency Exercise cached OpenAI and Foundry wrappers under overlapping streaming and non-streaming runs, including session and function-call isolation. Document the pinned behavior and track the upstream concurrency contract. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9d876897-6c4c-4e04-ac52-33be4d28676f --- README.md | 12 - docs/architecture.md | 10 +- docs/frds/0007-multi-agent-delegation.md | 6 +- tests/test_client_manager_concurrency.py | 300 +++++++++++++++++++++++ 4 files changed, 314 insertions(+), 14 deletions(-) create mode 100644 tests/test_client_manager_concurrency.py diff --git a/README.md b/README.md index 7b20797b..fbf8f517 100644 --- a/README.md +++ b/README.md @@ -56,18 +56,6 @@ 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. -### Advanced client manager lifecycle - -`ClientManager` is an advanced extension point used by test fakes, custom -MAF-compatible provider clients, and embedding hosts. Install a custom manager -once, before the default manager is requested. Replacing any active manager -requires `await shutdown_client_manager()` before `set_client_manager(...)`. - -Azure Functions does not currently expose a supported async worker-shutdown -hook, so the runtime keeps its default clients for the worker lifetime rather -than attempting cleanup from `atexit` or a signal handler. The host gap is tracked in -[azure-functions-python-worker#1904](https://github.com/Azure/azure-functions-python-worker/issues/1904). - ## Quick Start ### 1. Create the agent file diff --git a/docs/architecture.md b/docs/architecture.md index aa45c90f..b1ff495d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -343,7 +343,7 @@ 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 [Advanced client manager lifecycle](https://github.com/Azure/azure-functions-agents-runtime/blob/main/README.md#advanced-client-manager-lifecycle). +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. @@ -366,6 +366,14 @@ 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 diff --git a/docs/frds/0007-multi-agent-delegation.md b/docs/frds/0007-multi-agent-delegation.md index 0f5ecb5f..0581f127 100644 --- a/docs/frds/0007-multi-agent-delegation.md +++ b/docs/frds/0007-multi-agent-delegation.md @@ -643,6 +643,7 @@ handoff participant may itself declare `subagents` and delegate. | 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 @@ -680,7 +681,10 @@ handoff participant may itself declare `subagents` and delegate. 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). + (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/tests/test_client_manager_concurrency.py b/tests/test_client_manager_concurrency.py new file mode 100644 index 00000000..1861267c --- /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() From c1bfdaefaf29e901f5daaca9ff2e6468128752cb Mon Sep 17 00:00:00 2001 From: Tsuyoshi Ushio Date: Thu, 13 Aug 2026 11:19:58 -0700 Subject: [PATCH 5/5] ci: rerun after transient Node timeout The equivalent public pipeline rerun (build 297944) passed on both Python 3.13 and 3.14. This empty commit refreshes the PR check that cannot be re-requested through the Azure Pipelines GitHub app. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9d876897-6c4c-4e04-ac52-33be4d28676f