Repository navigation
Python: Add async context managers and ownership close to embedding clients - #9815
Open
Wilson_Jiang (Wilson0511) wants to merge 2 commits into
Open
Wilson_Jiang (Wilson0511) wants to merge 2 commits into
Wilson_Jiang (Wilson0511) wants to merge 2 commits into
Conversation
…lients microsoft#9748 gave chat clients deterministic lifecycle (context enter/exit, close, connector-owned resource release) but left the embedding client family without any of it: BaseEmbeddingClient has no __aenter__/ __aexit__/close, and the OpenAI, Bedrock, and Ollama embedding clients never release the SDK clients they create internally. A client built with default settings owns an AsyncOpenAI httpx pool (or a Bedrock runtime client, or an Ollama httpx transport) that only GC ever closes, and callers cannot use 'async with' at all. The Gemini embedding client already implemented close(); this brings the rest of the family to the same contract: - BaseEmbeddingClient gains the same context trio as BaseChatClient, with a no-op default close() and sync-close tolerance. - OpenAI, Bedrock, and Ollama embedding clients track _owns_client and close only the SDK client they created, leaving caller-supplied clients/credentials open; idempotent via _client_closed. Ollama closes the underlying httpx transport because ollama<0.6 has no AsyncClient aclose. 🤖 Generated with ZCode Co-Authored-By: ZCode <noreply@zcode.dev>
Wilson_Jiang (Wilson0511)
temporarily deployed
to
github-app-auth
October 29, 2026 03:11 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
Wilson_Jiang (Wilson0511)
requested review from
Darren Cohen (dargilco),
Jose Alvarez (jpalvarezl) and
sophia-ramsey
as code owners
October 29, 2026 03:11
Wilson_Jiang (Wilson0511)
temporarily deployed
to
github-app-auth
October 29, 2026 03:11 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
agent-framework-automation
Bot
added
the
python
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Embedding lifecycle tests do not verify cleanup call counts for owned versus injected clients or cover legacy synchronous close overrides on normal and exceptional exits.
2 open issues
What changed in this PR
Extends the deterministic lifecycle introduced for chat clients in #9748 to Python embedding clients, addressing #9814 while preserving caller-owned resources.
Changes:
- Adds async context management and a default no-op
close()toBaseEmbeddingClient. - Adds ownership-aware, idempotent cleanup for OpenAI, Bedrock, and Ollama embedding clients.
- Adds regression tests for lifecycle behavior, resource ownership, and serialization.
| File | Description |
|---|---|
python/packages/openai/tests/openai/test_openai_embedding_client.py |
Tests owned versus injected SDK-client cleanup. |
python/packages/openai/agent_framework_openai/_embedding_client.py |
Tracks ownership and closes internally created SDK clients. |
python/packages/ollama/tests/ollama/test_ollama_embedding_client.py |
Tests transport cleanup and injected-client preservation. |
python/packages/ollama/agent_framework_ollama/_embedding_client.py |
Adds ownership-aware HTTP transport cleanup. |
python/packages/core/tests/core/test_clients.py |
Tests default embedding lifecycle and context-exit cleanup. |
python/packages/core/agent_framework/_clients.py |
Defines the shared embedding-client lifecycle contract. |
python/packages/bedrock/tests/test_bedrock_client.py |
Tests runtime-client ownership and repeated cleanup. |
python/packages/bedrock/agent_framework_bedrock/_embedding_client.py |
Closes owned runtime clients off the event loop. |
Review effort: Balanced
Comment on lines
+1466
to
+1478
| @pytest.mark.parametrize("error", [None, RuntimeError("body failed")]) | ||
| async def test_base_embedding_client_context_calls_close(error: BaseException | None) -> None: | ||
| client = _MockEmbeddingClient() | ||
| with patch.object(client, "close", new_callable=AsyncMock) as close: | ||
| if error is None: | ||
| async with client: | ||
| close.assert_not_awaited() | ||
| else: | ||
| with pytest.raises(type(error)) as raised: | ||
| async with client: | ||
| raise error | ||
| assert raised.value is error | ||
| close.assert_awaited_once_with() |
There was a problem hiding this comment.
Added test_base_embedding_client_context_accepts_sync_close in the follow-up commit, mirroring the chat-client test (sync close() returning None, both normal and exceptional exits).
Comment on lines
+198
to
+201
| async with client as entered: | ||
| assert entered is client | ||
| await client.close() | ||
| assert client.client._client.is_closed is not injected |
There was a problem hiding this comment.
Fixed in the follow-up commit: the Ollama test now wraps the transport's aclose() and asserts await_count == (0 if injected else 1), so removing the _client_closed guard would fail it; is_closed is kept as a secondary state assertion.
Review follow-up on microsoft#9815: add the synchronous-close context-exit test for BaseEmbeddingClient (both normal and exceptional exits), and make the Ollama ownership test assert aclose() call counts so the idempotency guard is actually pinned (is_closed alone stays true on double close). 🤖 Generated with ZCode Co-Authored-By: ZCode <noreply@zcode.dev>
Wilson_Jiang (Wilson0511)
had a problem deploying
to
github-app-auth
October 29, 2026 03:32 — with
GitHub Actions
Failure
GitHub Actions
Failure
Member
|
Please fix the merge conflict Wilson_Jiang (@Wilson0511) |
This branch had an error being deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #9814
Description
#9748 gave chat clients deterministic lifecycle but left the embedding client family without any of it:
BaseEmbeddingClienthas no context trio orclose(), and the OpenAI, Bedrock, and Ollama embedding clients never release the SDK clients they create internally (httpx pool / bedrock-runtime client / httpx transport) — only GC ever closes them. The Gemini embedding client already implementsclose(); this brings the rest of the family to the same contract:BaseEmbeddingClientgains the same__aenter__/__aexit__/close()trio Python: Add async context managers to chat clients #9748 added toBaseChatClient: no-op default close, caller-usable without a context, syncclose()overrides still accepted on exit.RawOpenAIEmbeddingClient/RawBedrockEmbeddingClient/RawOllamaEmbeddingClienttrack_owns_client(true only when the SDK client was created internally) andclose()releases exactly that client, guarded by_client_closedfor idempotency. Caller-supplied clients, sessions, and credentials stay open. Bedrock closes its sync boto3 client viaasyncio.to_thread, mirroring Python: Add async context managers to chat clients #9748's chat client. Ollama closes the underlying httpx transport becauseollama<0.6has noAsyncClient.aclose(the constructor already reads the same private attribute).close()now also getsasync withsupport from the base trio.RegressionCheck
mainand pass here:test_base_embedding_client_default_lifecycle,test_base_embedding_client_context_calls_close(core);test_openai_embedding_client_context_ownership(openai);test_bedrock_embedding_client_context_ownership(bedrock);test_ollama_embedding_client_context_ownership(ollama). Ownership tests parametrize injected vs internally-created SDK clients and assert: context exit + explicitclose()close only the owned client (await count 1 vs 0), the SDK client ends closed only when owned, and_owns_client/_client_closeddo not leak intoto_dict().tests/for core, openai, bedrock, ollama — failure sets byte-identical tomainat 4a272d2 (pre-existing environment failures only; zero new). openai/bedrock/ollama suites fully green including the new tests.🤖 Generated with ZCode
Co-Authored-By: ZCode noreply@zcode.dev