From 2b251d904f1cb045ec3b7e13411cac3c1aa1e019 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Wed, 15 Apr 2026 19:30:12 +0200 Subject: [PATCH 1/9] Python: Fix reasoning replay when store=False (#5250) * fix reasoning content when store=False * Remove accidental worktree entries Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * remove local session sample * removed left over files * Add attribution override regression test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agent_framework_openai/_chat_client.py | 95 ++++--- .../tests/openai/test_openai_chat_client.py | 236 ++++++++++++++++++ 2 files changed, 301 insertions(+), 30 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index ecd17c4b5e..0f66974e49 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -1161,7 +1161,16 @@ class RawOpenAIChatClient( # type: ignore[misc] # First turn: prepend instructions as system message messages = prepend_instructions_to_messages(list(messages), instructions, role="system") # Continuation turn: instructions already exist in conversation context, skip prepending - request_input = self._prepare_messages_for_openai(messages) + request_uses_service_side_storage = False + for key in ("conversation_id", "previous_response_id", "conversation"): + value = options.get(key) + if isinstance(value, str) and value: + request_uses_service_side_storage = True + break + request_input = self._prepare_messages_for_openai( + messages, + request_uses_service_side_storage=request_uses_service_side_storage, + ) if not request_input: raise ChatClientInvalidRequestException("Messages are required for chat completions") conversation_id = options.get("conversation_id") @@ -1235,7 +1244,12 @@ class RawOpenAIChatClient( # type: ignore[misc] raise ValueError("model must be a non-empty string") options["model"] = self.model - def _prepare_messages_for_openai(self, chat_messages: Sequence[Message]) -> list[dict[str, Any]]: + def _prepare_messages_for_openai( + self, + chat_messages: Sequence[Message], + *, + request_uses_service_side_storage: bool = True, + ) -> list[dict[str, Any]]: """Prepare the chat messages for a request. Allowing customization of the key names for role/author, and optionally overriding the role. @@ -1248,31 +1262,27 @@ class RawOpenAIChatClient( # type: ignore[misc] Args: chat_messages: The chat history to prepare. + request_uses_service_side_storage: Whether this request continues a service-managed + response/conversation and can safely reference service-scoped response items. Returns: The prepared chat messages for a request. """ - list_of_list = [self._prepare_message_for_openai(message) for message in chat_messages] + list_of_list = [ + self._prepare_message_for_openai( + message, + request_uses_service_side_storage=request_uses_service_side_storage, + ) + for message in chat_messages + ] # Flatten the list of lists into a single list return list(chain.from_iterable(list_of_list)) - @staticmethod - def _message_replays_provider_context(message: Message) -> bool: - """Return whether the message came from provider-attributed replay context. - - Responses ``fc_id`` values are response-scoped and only valid while replaying - the same live tool loop. Once a message comes back through a context provider - (for example, loaded session history), that message is historical input and - must not reuse the original response-scoped ``fc_id``. - """ - additional_properties = getattr(message, "additional_properties", None) - if not additional_properties: - return False - return "_attribution" in additional_properties - def _prepare_message_for_openai( self, message: Message, + *, + request_uses_service_side_storage: bool = True, ) -> list[dict[str, Any]]: """Prepare a chat message for the OpenAI Responses API format.""" all_messages: list[dict[str, Any]] = [] @@ -1280,34 +1290,63 @@ class RawOpenAIChatClient( # type: ignore[misc] "type": "message", "role": message.role, } + additional_properties = message.additional_properties + replays_local_storage = "_attribution" in additional_properties + uses_service_side_storage = request_uses_service_side_storage and not replays_local_storage # Reasoning items are only valid in input when they directly preceded a function_call - # in the same response. Including a reasoning item that preceded a text response + # in the same response. Including a reasoning item that preceded a text response # (i.e. no function_call in the same message) causes an API error: # "reasoning was provided without its required following item." + # + # Local storage is stricter: response-scoped reasoning items (rs_*) cannot be replayed + # back to the service unless that message is using service-side storage. + # In that mode we omit reasoning items and rely on function call + tool output replay. has_function_call = any(c.type == "function_call" for c in message.contents) for content in message.contents: match content.type: case "text_reasoning": - if not has_function_call: + if not uses_service_side_storage or not has_function_call: continue # reasoning not followed by a function_call is invalid in input - reasoning = self._prepare_content_for_openai(message.role, content, message=message) + reasoning = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) if reasoning: all_messages.append(reasoning) case "function_result": new_args: dict[str, Any] = {} - new_args.update(self._prepare_content_for_openai(message.role, content, message=message)) + new_args.update( + self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) + ) if new_args: all_messages.append(new_args) case "function_call": - function_call = self._prepare_content_for_openai(message.role, content, message=message) + function_call = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) if function_call: all_messages.append(function_call) case "function_approval_response" | "function_approval_request": - prepared = self._prepare_content_for_openai(message.role, content, message=message) + prepared = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) if prepared: all_messages.append(prepared) case _: - prepared_content = self._prepare_content_for_openai(message.role, content, message=message) + prepared_content = self._prepare_content_for_openai( + message.role, + content, + replays_local_storage=replays_local_storage, + ) if prepared_content: if "content" not in args: args["content"] = [] @@ -1321,7 +1360,7 @@ class RawOpenAIChatClient( # type: ignore[misc] role: Role | str, content: Content, *, - message: Message | None = None, + replays_local_storage: bool = False, ) -> dict[str, Any]: """Prepare content for the OpenAI Responses API format.""" role = Role(role) @@ -1401,11 +1440,7 @@ class RawOpenAIChatClient( # type: ignore[misc] logger.warning(f"FunctionCallContent missing call_id for function '{content.name}'") return {} fc_id = content.call_id - if ( - message is not None - and not self._message_replays_provider_context(message) - and content.additional_properties - ): + if not replays_local_storage and content.additional_properties: live_fc_id = content.additional_properties.get("fc_id") if isinstance(live_fc_id, str) and live_fc_id: fc_id = live_fc_id diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 1e18f85273..fe4ee4124b 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -1015,6 +1015,84 @@ async def test_shell_call_is_invoked_as_local_shell_function_loop() -> None: assert len(local_shell_outputs) == 0 +async def test_tool_loop_store_false_omits_reasoning_items_from_second_request() -> None: + """Stateless tool-loop replay must omit response-scoped reasoning items.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + mock_response1 = MagicMock() + mock_response1.output_parsed = None + mock_response1.metadata = {} + mock_response1.usage = None + mock_response1.id = "resp-1" + mock_response1.model = "test-model" + mock_response1.created_at = 1000000000 + mock_response1.status = "completed" + mock_response1.finish_reason = "tool_calls" + mock_response1.incomplete = None + mock_response1.conversation = None + + mock_reasoning_item = MagicMock() + mock_reasoning_item.type = "reasoning" + mock_reasoning_item.id = "rs_local_only" + mock_reasoning_item.content = [] + mock_reasoning_item.summary = [] + mock_reasoning_item.encrypted_content = None + + mock_function_call_item = MagicMock() + mock_function_call_item.type = "function_call" + mock_function_call_item.id = "fc_tool123" + mock_function_call_item.call_id = "call_123" + mock_function_call_item.name = "get_weather" + mock_function_call_item.arguments = '{"location":"Amsterdam"}' + mock_function_call_item.status = "completed" + + mock_response1.output = [mock_reasoning_item, mock_function_call_item] + + mock_response2 = MagicMock() + mock_response2.output_parsed = None + mock_response2.metadata = {} + mock_response2.usage = None + mock_response2.id = "resp-2" + mock_response2.model = "test-model" + mock_response2.created_at = 1000000001 + mock_response2.status = "completed" + mock_response2.finish_reason = "stop" + mock_response2.incomplete = None + mock_response2.conversation = None + + mock_text_item = MagicMock() + mock_text_item.type = "message" + mock_text_content = MagicMock() + mock_text_content.type = "output_text" + mock_text_content.text = "The weather in Amsterdam is sunny." + mock_text_item.content = [mock_text_content] + mock_response2.output = [mock_text_item] + + with patch.object(client.client.responses, "create", side_effect=[mock_response1, mock_response2]) as mock_create: + response = await client.get_response( + messages=[Message(role="user", contents=["What's the weather in Amsterdam?"])], + options={ + "store": False, + "tools": [get_weather], + "tool_choice": {"mode": "required", "required_function_name": "get_weather"}, + }, + ) + + assert response.text == "The weather in Amsterdam is sunny." + assert mock_create.call_count == 2 + + second_call_input = mock_create.call_args_list[1].kwargs["input"] + assert not any(item.get("type") == "reasoning" for item in second_call_input) + + function_calls = [item for item in second_call_input if item.get("type") == "function_call"] + assert len(function_calls) == 1 + assert function_calls[0]["id"] == "fc_tool123" + + function_outputs = [item for item in second_call_input if item.get("type") == "function_call_output"] + assert len(function_outputs) == 1 + assert function_outputs[0]["call_id"] == "call_123" + + def test_response_content_creation_with_shell_call() -> None: """Test _parse_response_from_openai with shell_call output.""" client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -3221,6 +3299,164 @@ async def test_prepare_options_store_parameter_handling() -> None: assert "previous_response_id" not in options +async def test_prepare_options_store_false_omits_reasoning_items_for_stateless_replay() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + messages = [ + Message(role="user", contents=[Content.from_text(text="search for hotels")]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_test123", + text="I need to search for hotels", + additional_properties={"status": "completed"}, + ), + Content.from_function_call( + call_id="call_1", + name="search_hotels", + arguments='{"city": "Paris"}', + additional_properties={"fc_id": "fc_test456"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="call_1", + result="Found 3 hotels in Paris", + ), + ], + ), + ] + + options = await client._prepare_options(messages, ChatOptions(store=False)) # type: ignore[arg-type] + + assert not any(item.get("type") == "reasoning" for item in options["input"]) + assert any(item.get("type") == "function_call" for item in options["input"]) + assert any(item.get("type") == "function_call_output" for item in options["input"]) + + +async def test_prepare_options_with_conversation_id_keeps_reasoning_items() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + messages = [ + Message(role="user", contents=[Content.from_text(text="search for hotels")]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_test123", + text="I need to search for hotels", + additional_properties={"status": "completed"}, + ), + Content.from_function_call( + call_id="call_1", + name="search_hotels", + arguments='{"city": "Paris"}', + additional_properties={"fc_id": "fc_test456"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="call_1", + result="Found 3 hotels in Paris", + ), + ], + ), + ] + + options = await client._prepare_options( + messages, + ChatOptions(store=False, conversation_id="resp_prev123"), # type: ignore[arg-type] + ) + + reasoning_items = [item for item in options["input"] if item.get("type") == "reasoning"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["id"] == "rs_test123" + assert options["previous_response_id"] == "resp_prev123" + + +async def test_prepare_options_with_conversation_id_omits_reasoning_items_for_attributed_replay() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + messages = [ + Message(role="user", contents=[Content.from_text(text="search for hotels")]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_history123", + text="I need to search history for hotels", + additional_properties={"status": "completed"}, + ), + Content.from_function_call( + call_id="call_history", + name="search_hotels", + arguments='{"city": "Paris"}', + additional_properties={"fc_id": "fc_history456"}, + ), + ], + additional_properties={"_attribution": {"source_id": "history", "source_type": "InMemoryHistoryProvider"}}, + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="call_history", + result="Found 3 hotels in Paris", + ), + ], + ), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_live123", + text="I should refine the search for a live follow-up", + additional_properties={"status": "completed"}, + ), + Content.from_function_call( + call_id="call_live", + name="search_hotels", + arguments='{"city": "London"}', + additional_properties={"fc_id": "fc_live456"}, + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="call_live", + result="Found 4 hotels in London", + ), + ], + ), + ] + + options = await client._prepare_options( + messages, + ChatOptions(store=False, conversation_id="resp_prev123"), # type: ignore[arg-type] + ) + + reasoning_items = [item for item in options["input"] if item.get("type") == "reasoning"] + assert [item["id"] for item in reasoning_items] == ["rs_live123"] + assert any( + item.get("type") == "function_call" and item.get("call_id") == "call_history" for item in options["input"] + ) + assert any(item.get("type") == "function_call" and item.get("call_id") == "call_live" for item in options["input"]) + assert any( + item.get("type") == "function_call_output" and item.get("call_id") == "call_history" + for item in options["input"] + ) + assert any( + item.get("type") == "function_call_output" and item.get("call_id") == "call_live" for item in options["input"] + ) + assert options["previous_response_id"] == "resp_prev123" + + def _create_mock_responses_text_response(*, response_id: str) -> MagicMock: mock_response = MagicMock() mock_response.id = response_id From 68b93641b6802abd2d00a5191cb16c6074d39fe1 Mon Sep 17 00:00:00 2001 From: Giles Odigwe <79032838+giles17@users.noreply.github.com> Date: Wed, 15 Apr 2026 11:22:15 -0700 Subject: [PATCH 2/9] Python: Bump agent-framework-devui to 1.0.0b260414 for release (#5259) Update devui version and changelog for the streaming memory fix release. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/CHANGELOG.md | 5 +++++ python/packages/devui/pyproject.toml | 2 +- python/uv.lock | 8 +++++++- 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 0ae0df1454..4adafae53c 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -10,6 +10,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - **agent-framework-azure-cosmos**: [BREAKING] `CosmosCheckpointStorage` now uses restricted pickle deserialization by default, matching `FileCheckpointStorage` behavior. If your checkpoints contain application-defined types, pass them via `allowed_checkpoint_types=["my_app.models:MyState"]`. ([#5200](https://github.com/microsoft/agent-framework/issues/5200)) +## [devui-1.0.0b260414] - 2026-04-14 + +### Fixed +- **agent-framework-devui**: Fix streaming memory growth in DevUI frontend ([#5221](https://github.com/microsoft/agent-framework/pull/5221)) + ## [1.0.1] - 2026-04-09 ### Added diff --git a/python/packages/devui/pyproject.toml b/python/packages/devui/pyproject.toml index b12fbef49b..65a595eb48 100644 --- a/python/packages/devui/pyproject.toml +++ b/python/packages/devui/pyproject.toml @@ -4,7 +4,7 @@ description = "Debug UI for Microsoft Agent Framework with OpenAI-compatible API authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] readme = "README.md" requires-python = ">=3.10" -version = "1.0.0b260409" +version = "1.0.0b260414" license-files = ["LICENSE"] urls.homepage = "https://github.com/microsoft/agent-framework" urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" diff --git a/python/uv.lock b/python/uv.lock index 138c6b49fb..7755412101 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -416,7 +416,7 @@ dev = [{ name = "types-pyyaml", specifier = "==6.0.12.20250915" }] [[package]] name = "agent-framework-devui" -version = "1.0.0b260409" +version = "1.0.0b260414" source = { editable = "packages/devui" } dependencies = [ { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -2406,6 +2406,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/3f/9859f655d11901e7b2996c6e3d33e0caa9a1d4572c3bc61ed0faa64b2f4c/greenlet-3.3.2-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:9bc885b89709d901859cf95179ec9f6bb67a3d2bb1f0e88456461bd4b7f8fd0d", size = 277747, upload-time = "2026-02-20T20:16:21.325Z" }, { url = "https://files.pythonhosted.org/packages/fb/07/cb284a8b5c6498dbd7cba35d31380bb123d7dceaa7907f606c8ff5993cbf/greenlet-3.3.2-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b568183cf65b94919be4438dc28416b234b678c608cafac8874dfeeb2a9bbe13", size = 579202, upload-time = "2026-02-20T20:47:28.955Z" }, { url = "https://files.pythonhosted.org/packages/ed/45/67922992b3a152f726163b19f890a85129a992f39607a2a53155de3448b8/greenlet-3.3.2-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:527fec58dc9f90efd594b9b700662ed3fb2493c2122067ac9c740d98080a620e", size = 590620, upload-time = "2026-02-20T20:55:55.581Z" }, + { url = "https://files.pythonhosted.org/packages/03/5f/6e2a7d80c353587751ef3d44bb947f0565ec008a2e0927821c007e96d3a7/greenlet-3.3.2-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:508c7f01f1791fbc8e011bd508f6794cb95397fdb198a46cb6635eb5b78d85a7", size = 602132, upload-time = "2026-02-20T21:02:43.261Z" }, { url = "https://files.pythonhosted.org/packages/ad/55/9f1ebb5a825215fadcc0f7d5073f6e79e3007e3282b14b22d6aba7ca6cb8/greenlet-3.3.2-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ad0c8917dd42a819fe77e6bdfcb84e3379c0de956469301d9fd36427a1ca501f", size = 591729, upload-time = "2026-02-20T20:20:58.395Z" }, { url = "https://files.pythonhosted.org/packages/24/b4/21f5455773d37f94b866eb3cf5caed88d6cea6dd2c6e1f9c34f463cba3ec/greenlet-3.3.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:97245cc10e5515dbc8c3104b2928f7f02b6813002770cfaffaf9a6e0fc2b94ef", size = 1551946, upload-time = "2026-02-20T20:49:31.102Z" }, { url = "https://files.pythonhosted.org/packages/00/68/91f061a926abead128fe1a87f0b453ccf07368666bd59ffa46016627a930/greenlet-3.3.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8c1fdd7d1b309ff0da81d60a9688a8bd044ac4e18b250320a96fc68d31c209ca", size = 1618494, upload-time = "2026-02-20T20:21:06.541Z" }, @@ -2413,6 +2414,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f3/47/16400cb42d18d7a6bb46f0626852c1718612e35dcb0dffa16bbaffdf5dd2/greenlet-3.3.2-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:c56692189a7d1c7606cb794be0a8381470d95c57ce5be03fb3d0ef57c7853b86", size = 278890, upload-time = "2026-02-20T20:19:39.263Z" }, { url = "https://files.pythonhosted.org/packages/a3/90/42762b77a5b6aa96cd8c0e80612663d39211e8ae8a6cd47c7f1249a66262/greenlet-3.3.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ebd458fa8285960f382841da585e02201b53a5ec2bac6b156fc623b5ce4499f", size = 581120, upload-time = "2026-02-20T20:47:30.161Z" }, { url = "https://files.pythonhosted.org/packages/bf/6f/f3d64f4fa0a9c7b5c5b3c810ff1df614540d5aa7d519261b53fba55d4df9/greenlet-3.3.2-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a443358b33c4ec7b05b79a7c8b466f5d275025e750298be7340f8fc63dff2a55", size = 594363, upload-time = "2026-02-20T20:55:56.965Z" }, + { url = "https://files.pythonhosted.org/packages/9c/8b/1430a04657735a3f23116c2e0d5eb10220928846e4537a938a41b350bed6/greenlet-3.3.2-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4375a58e49522698d3e70cc0b801c19433021b5c37686f7ce9c65b0d5c8677d2", size = 605046, upload-time = "2026-02-20T21:02:45.234Z" }, { url = "https://files.pythonhosted.org/packages/72/83/3e06a52aca8128bdd4dcd67e932b809e76a96ab8c232a8b025b2850264c5/greenlet-3.3.2-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e2cd90d413acbf5e77ae41e5d3c9b3ac1d011a756d7284d7f3f2b806bbd6358", size = 594156, upload-time = "2026-02-20T20:20:59.955Z" }, { url = "https://files.pythonhosted.org/packages/70/79/0de5e62b873e08fe3cef7dbe84e5c4bc0e8ed0c7ff131bccb8405cd107c8/greenlet-3.3.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:442b6057453c8cb29b4fb36a2ac689382fc71112273726e2423f7f17dc73bf99", size = 1554649, upload-time = "2026-02-20T20:49:32.293Z" }, { url = "https://files.pythonhosted.org/packages/5a/00/32d30dee8389dc36d42170a9c66217757289e2afb0de59a3565260f38373/greenlet-3.3.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:45abe8eb6339518180d5a7fa47fa01945414d7cca5ecb745346fc6a87d2750be", size = 1619472, upload-time = "2026-02-20T20:21:07.966Z" }, @@ -2421,6 +2423,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/ab/1608e5a7578e62113506740b88066bf09888322a311cff602105e619bd87/greenlet-3.3.2-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:ac8d61d4343b799d1e526db579833d72f23759c71e07181c2d2944e429eb09cd", size = 280358, upload-time = "2026-02-20T20:17:43.971Z" }, { url = "https://files.pythonhosted.org/packages/a5/23/0eae412a4ade4e6623ff7626e38998cb9b11e9ff1ebacaa021e4e108ec15/greenlet-3.3.2-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ceec72030dae6ac0c8ed7591b96b70410a8be370b6a477b1dbc072856ad02bd", size = 601217, upload-time = "2026-02-20T20:47:31.462Z" }, { url = "https://files.pythonhosted.org/packages/f8/16/5b1678a9c07098ecb9ab2dd159fafaf12e963293e61ee8d10ecb55273e5e/greenlet-3.3.2-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a2a5be83a45ce6188c045bcc44b0ee037d6a518978de9a5d97438548b953a1ac", size = 611792, upload-time = "2026-02-20T20:55:58.423Z" }, + { url = "https://files.pythonhosted.org/packages/5c/c5/cc09412a29e43406eba18d61c70baa936e299bc27e074e2be3806ed29098/greenlet-3.3.2-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ae9e21c84035c490506c17002f5c8ab25f980205c3e61ddb3a2a2a2e6c411fcb", size = 626250, upload-time = "2026-02-20T21:02:46.596Z" }, { url = "https://files.pythonhosted.org/packages/50/1f/5155f55bd71cabd03765a4aac9ac446be129895271f73872c36ebd4b04b6/greenlet-3.3.2-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:43e99d1749147ac21dde49b99c9abffcbc1e2d55c67501465ef0930d6e78e070", size = 613875, upload-time = "2026-02-20T20:21:01.102Z" }, { url = "https://files.pythonhosted.org/packages/fc/dd/845f249c3fcd69e32df80cdab059b4be8b766ef5830a3d0aa9d6cad55beb/greenlet-3.3.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4c956a19350e2c37f2c48b336a3afb4bff120b36076d9d7fb68cb44e05d95b79", size = 1571467, upload-time = "2026-02-20T20:49:33.495Z" }, { url = "https://files.pythonhosted.org/packages/2a/50/2649fe21fcc2b56659a452868e695634722a6655ba245d9f77f5656010bf/greenlet-3.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6c6f8ba97d17a1e7d664151284cb3315fc5f8353e75221ed4324f84eb162b395", size = 1640001, upload-time = "2026-02-20T20:21:09.154Z" }, @@ -2429,6 +2432,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ac/48/f8b875fa7dea7dd9b33245e37f065af59df6a25af2f9561efa8d822fde51/greenlet-3.3.2-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:aa6ac98bdfd716a749b84d4034486863fd81c3abde9aa3cf8eff9127981a4ae4", size = 279120, upload-time = "2026-02-20T20:19:01.9Z" }, { url = "https://files.pythonhosted.org/packages/49/8d/9771d03e7a8b1ee456511961e1b97a6d77ae1dea4a34a5b98eee706689d3/greenlet-3.3.2-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ab0c7e7901a00bc0a7284907273dc165b32e0d109a6713babd04471327ff7986", size = 603238, upload-time = "2026-02-20T20:47:32.873Z" }, { url = "https://files.pythonhosted.org/packages/59/0e/4223c2bbb63cd5c97f28ffb2a8aee71bdfb30b323c35d409450f51b91e3e/greenlet-3.3.2-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d248d8c23c67d2291ffd47af766e2a3aa9fa1c6703155c099feb11f526c63a92", size = 614219, upload-time = "2026-02-20T20:55:59.817Z" }, + { url = "https://files.pythonhosted.org/packages/94/2b/4d012a69759ac9d77210b8bfb128bc621125f5b20fc398bce3940d036b1c/greenlet-3.3.2-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ccd21bb86944ca9be6d967cf7691e658e43417782bce90b5d2faeda0ff78a7dd", size = 628268, upload-time = "2026-02-20T21:02:48.024Z" }, { url = "https://files.pythonhosted.org/packages/7a/34/259b28ea7a2a0c904b11cd36c79b8cef8019b26ee5dbe24e73b469dea347/greenlet-3.3.2-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6997d360a4e6a4e936c0f9625b1c20416b8a0ea18a8e19cabbefc712e7397ab", size = 616774, upload-time = "2026-02-20T20:21:02.454Z" }, { url = "https://files.pythonhosted.org/packages/0a/03/996c2d1689d486a6e199cb0f1cf9e4aa940c500e01bdf201299d7d61fa69/greenlet-3.3.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:64970c33a50551c7c50491671265d8954046cb6e8e2999aacdd60e439b70418a", size = 1571277, upload-time = "2026-02-20T20:49:34.795Z" }, { url = "https://files.pythonhosted.org/packages/d9/c4/2570fc07f34a39f2caf0bf9f24b0a1a0a47bc2e8e465b2c2424821389dfc/greenlet-3.3.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1a9172f5bf6bd88e6ba5a84e0a68afeac9dc7b6b412b245dd64f52d83c81e55b", size = 1640455, upload-time = "2026-02-20T20:21:10.261Z" }, @@ -2437,6 +2441,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/ae/8bffcbd373b57a5992cd077cbe8858fff39110480a9d50697091faea6f39/greenlet-3.3.2-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:8d1658d7291f9859beed69a776c10822a0a799bc4bfe1bd4272bb60e62507dab", size = 279650, upload-time = "2026-02-20T20:18:00.783Z" }, { url = "https://files.pythonhosted.org/packages/d1/c0/45f93f348fa49abf32ac8439938726c480bd96b2a3c6f4d949ec0124b69f/greenlet-3.3.2-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18cb1b7337bca281915b3c5d5ae19f4e76d35e1df80f4ad3c1a7be91fadf1082", size = 650295, upload-time = "2026-02-20T20:47:34.036Z" }, { url = "https://files.pythonhosted.org/packages/b3/de/dd7589b3f2b8372069ab3e4763ea5329940fc7ad9dcd3e272a37516d7c9b/greenlet-3.3.2-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c2e47408e8ce1c6f1ceea0dffcdf6ebb85cc09e55c7af407c99f1112016e45e9", size = 662163, upload-time = "2026-02-20T20:56:01.295Z" }, + { url = "https://files.pythonhosted.org/packages/cd/ac/85804f74f1ccea31ba518dcc8ee6f14c79f73fe36fa1beba38930806df09/greenlet-3.3.2-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e3cb43ce200f59483eb82949bf1835a99cf43d7571e900d7c8d5c62cdf25d2f9", size = 675371, upload-time = "2026-02-20T21:02:49.664Z" }, { url = "https://files.pythonhosted.org/packages/d2/d8/09bfa816572a4d83bccd6750df1926f79158b1c36c5f73786e26dbe4ee38/greenlet-3.3.2-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:63d10328839d1973e5ba35e98cccbca71b232b14051fd957b6f8b6e8e80d0506", size = 664160, upload-time = "2026-02-20T20:21:04.015Z" }, { url = "https://files.pythonhosted.org/packages/48/cf/56832f0c8255d27f6c35d41b5ec91168d74ec721d85f01a12131eec6b93c/greenlet-3.3.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8e4ab3cfb02993c8cc248ea73d7dae6cec0253e9afa311c9b37e603ca9fad2ce", size = 1619181, upload-time = "2026-02-20T20:49:36.052Z" }, { url = "https://files.pythonhosted.org/packages/0a/23/b90b60a4aabb4cec0796e55f25ffbfb579a907c3898cd2905c8918acaa16/greenlet-3.3.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:94ad81f0fd3c0c0681a018a976e5c2bd2ca2d9d94895f23e7bb1af4e8af4e2d5", size = 1687713, upload-time = "2026-02-20T20:21:11.684Z" }, @@ -2445,6 +2450,7 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/6d/8f2ef704e614bcf58ed43cfb8d87afa1c285e98194ab2cfad351bf04f81e/greenlet-3.3.2-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:e26e72bec7ab387ac80caa7496e0f908ff954f31065b0ffc1f8ecb1338b11b54", size = 286617, upload-time = "2026-02-20T20:19:29.856Z" }, { url = "https://files.pythonhosted.org/packages/5e/0d/93894161d307c6ea237a43988f27eba0947b360b99ac5239ad3fe09f0b47/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b466dff7a4ffda6ca975979bab80bdadde979e29fc947ac3be4451428d8b0e4", size = 655189, upload-time = "2026-02-20T20:47:35.742Z" }, { url = "https://files.pythonhosted.org/packages/f5/2c/d2d506ebd8abcb57386ec4f7ba20f4030cbe56eae541bc6fd6ef399c0b41/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b8bddc5b73c9720bea487b3bffdb1840fe4e3656fba3bd40aa1489e9f37877ff", size = 658225, upload-time = "2026-02-20T20:56:02.527Z" }, + { url = "https://files.pythonhosted.org/packages/d1/67/8197b7e7e602150938049d8e7f30de1660cfb87e4c8ee349b42b67bdb2e1/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:59b3e2c40f6706b05a9cd299c836c6aa2378cabe25d021acd80f13abf81181cf", size = 666581, upload-time = "2026-02-20T21:02:51.526Z" }, { url = "https://files.pythonhosted.org/packages/8e/30/3a09155fbf728673a1dea713572d2d31159f824a37c22da82127056c44e4/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b26b0f4428b871a751968285a1ac9648944cea09807177ac639b030bddebcea4", size = 657907, upload-time = "2026-02-20T20:21:05.259Z" }, { url = "https://files.pythonhosted.org/packages/f3/fd/d05a4b7acd0154ed758797f0a43b4c0962a843bedfe980115e842c5b2d08/greenlet-3.3.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1fb39a11ee2e4d94be9a76671482be9398560955c9e568550de0224e41104727", size = 1618857, upload-time = "2026-02-20T20:49:37.309Z" }, { url = "https://files.pythonhosted.org/packages/6f/e1/50ee92a5db521de8f35075b5eff060dd43d39ebd46c2181a2042f7070385/greenlet-3.3.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:20154044d9085151bc309e7689d6f7ba10027f8f5a8c0676ad398b951913d89e", size = 1680010, upload-time = "2026-02-20T20:21:13.427Z" }, From eab7f09d03387a2b393f9785963353c1a09e8b6b Mon Sep 17 00:00:00 2001 From: S3rj <31356555+Serjbory@users.noreply.github.com> Date: Thu, 16 Apr 2026 00:08:01 +0200 Subject: [PATCH 3/9] Forward provider config to SessionConfig in GitHubCopilotAgent (fixes #5190) (#5195) Co-authored-by: Sergey Borisov --- .../agent_framework_github_copilot/_agent.py | 13 +- .../tests/test_github_copilot_agent.py | 193 ++++++++++++++++++ 2 files changed, 205 insertions(+), 1 deletion(-) diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 4744013b56..8015f561d9 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -31,7 +31,7 @@ from agent_framework.exceptions import AgentException try: from copilot import CopilotClient, CopilotSession, SubprocessConfig from copilot.generated.session_events import PermissionRequest, SessionEvent, SessionEventType - from copilot.session import MCPServerConfig, PermissionRequestResult, SystemMessageConfig + from copilot.session import MCPServerConfig, PermissionRequestResult, ProviderConfig, SystemMessageConfig from copilot.tools import Tool as CopilotTool from copilot.tools import ToolInvocation, ToolResult except ImportError as _copilot_import_error: @@ -120,6 +120,12 @@ class GitHubCopilotOptions(TypedDict, total=False): Supports both local (stdio) and remote (HTTP/SSE) servers. """ + provider: ProviderConfig + """Custom API provider configuration for BYOK (Bring Your Own Key) scenarios. + Allows routing requests through your own OpenAI, Azure, or Anthropic endpoint + instead of the default GitHub Copilot backend. + """ + OptionsT = TypeVar( "OptionsT", @@ -232,6 +238,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): log_level = opts.pop("log_level", None) on_permission_request: PermissionHandlerType | None = opts.pop("on_permission_request", None) mcp_servers: dict[str, MCPServerConfig] | None = opts.pop("mcp_servers", None) + provider: ProviderConfig | None = opts.pop("provider", None) self._settings = load_settings( GitHubCopilotSettings, @@ -247,6 +254,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): self._tools = normalize_tools(tools) self._permission_handler = on_permission_request self._mcp_servers = mcp_servers + self._provider = provider self._default_options = opts self._started = False @@ -730,6 +738,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): opts.get("on_permission_request") or self._permission_handler or _deny_all_permissions ) mcp_servers = opts.get("mcp_servers") or self._mcp_servers or None + provider = opts.get("provider") or self._provider or None tools = self._prepare_tools(self._tools) if self._tools else None return await self._client.create_session( @@ -739,6 +748,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): system_message=system_message or None, tools=tools or None, mcp_servers=mcp_servers or None, + provider=provider or None, ) async def _resume_session(self, session_id: str, streaming: bool) -> CopilotSession: @@ -755,4 +765,5 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): streaming=streaming, tools=tools or None, mcp_servers=self._mcp_servers or None, + provider=self._provider or None, ) diff --git a/python/packages/github_copilot/tests/test_github_copilot_agent.py b/python/packages/github_copilot/tests/test_github_copilot_agent.py index 17e9432e40..9e1459db93 100644 --- a/python/packages/github_copilot/tests/test_github_copilot_agent.py +++ b/python/packages/github_copilot/tests/test_github_copilot_agent.py @@ -861,6 +861,7 @@ class TestGitHubCopilotAgentSessionManagement: streaming=unittest.mock.ANY, tools=unittest.mock.ANY, mcp_servers=unittest.mock.ANY, + provider=unittest.mock.ANY, ) async def test_session_config_includes_model( @@ -1084,6 +1085,198 @@ class TestGitHubCopilotAgentMCPServers: assert config["mcp_servers"] is None +class TestGitHubCopilotAgentProvider: + """Test cases for provider configuration (BYOK / Managed Identity).""" + + async def test_provider_passed_to_create_session( + self, + mock_client: MagicMock, + ) -> None: + """Test that provider config is passed through to create_session.""" + from copilot.session import ProviderConfig + + provider: ProviderConfig = { + "type": "azure", + "base_url": "https://my-resource.openai.azure.com", + "bearer_token": "test-token", + } + + agent: GitHubCopilotAgent[GitHubCopilotOptions] = GitHubCopilotAgent( + client=mock_client, + default_options={"provider": provider}, + ) + await agent.start() + + await agent._get_or_create_session(AgentSession()) # type: ignore + + call_args = mock_client.create_session.call_args + config = call_args.kwargs + assert config["provider"]["type"] == "azure" + assert config["provider"]["base_url"] == "https://my-resource.openai.azure.com" + assert config["provider"]["bearer_token"] == "test-token" + + async def test_provider_passed_to_resume_session( + self, + mock_client: MagicMock, + ) -> None: + """Test that provider config is passed through to resume_session.""" + from copilot.session import ProviderConfig + + provider: ProviderConfig = { + "type": "azure", + "base_url": "https://my-resource.openai.azure.com", + "bearer_token": "test-token", + } + + agent: GitHubCopilotAgent[GitHubCopilotOptions] = GitHubCopilotAgent( + client=mock_client, + default_options={"provider": provider}, + ) + await agent.start() + + session = AgentSession() + session.service_session_id = "existing-session-id" + + await agent._get_or_create_session(session) # type: ignore + + mock_client.resume_session.assert_called_once() + call_args = mock_client.resume_session.call_args + config = call_args.kwargs + assert config["provider"]["type"] == "azure" + + async def test_session_config_excludes_provider_when_not_set( + self, + mock_client: MagicMock, + ) -> None: + """Test that provider is None in session config when not set.""" + agent = GitHubCopilotAgent(client=mock_client) + await agent.start() + + await agent._get_or_create_session(AgentSession()) # type: ignore + + call_args = mock_client.create_session.call_args + config = call_args.kwargs + assert config["provider"] is None + + async def test_resume_session_excludes_provider_when_not_set( + self, + mock_client: MagicMock, + ) -> None: + """Test that provider is None in resume session config when not set.""" + agent = GitHubCopilotAgent(client=mock_client) + await agent.start() + + session = AgentSession() + session.service_session_id = "existing-session-id" + + await agent._get_or_create_session(session) # type: ignore + + call_args = mock_client.resume_session.call_args + config = call_args.kwargs + assert config["provider"] is None + + async def test_runtime_provider_takes_precedence( + self, + mock_client: MagicMock, + ) -> None: + """Test that runtime provider options override default_options provider.""" + from copilot.session import ProviderConfig + + default_provider: ProviderConfig = { + "type": "azure", + "base_url": "https://default.openai.azure.com", + "bearer_token": "default-token", + } + runtime_provider: ProviderConfig = { + "type": "openai", + "base_url": "https://runtime.openai.com", + "api_key": "runtime-key", + } + + agent: GitHubCopilotAgent[GitHubCopilotOptions] = GitHubCopilotAgent( + client=mock_client, + default_options={"provider": default_provider}, + ) + await agent.start() + + await agent._get_or_create_session( # type: ignore + AgentSession(), + runtime_options={"provider": runtime_provider}, + ) + + call_args = mock_client.create_session.call_args + config = call_args.kwargs + assert config["provider"]["type"] == "openai" + assert config["provider"]["base_url"] == "https://runtime.openai.com" + + async def test_provider_not_leaked_into_default_options( + self, + mock_client: MagicMock, + ) -> None: + """Test that provider is popped from opts and not left in _default_options.""" + from copilot.session import ProviderConfig + + provider: ProviderConfig = { + "type": "azure", + "base_url": "https://my-resource.openai.azure.com", + "bearer_token": "test-token", + } + + agent: GitHubCopilotAgent[GitHubCopilotOptions] = GitHubCopilotAgent( + client=mock_client, + default_options={"provider": provider, "model": "gpt-5"}, + ) + + assert "provider" not in agent._default_options + assert agent._provider is not None + assert agent._provider["type"] == "azure" + + async def test_provider_coexists_with_other_options( + self, + mock_client: MagicMock, + ) -> None: + """Test that provider works alongside model, tools, and mcp_servers.""" + from copilot.session import MCPServerConfig, ProviderConfig + + provider: ProviderConfig = { + "type": "azure", + "base_url": "https://my-resource.openai.azure.com", + "bearer_token": "test-token", + } + mcp_servers: dict[str, MCPServerConfig] = { + "test-server": { + "type": "stdio", + "command": "echo", + "args": ["hello"], + "tools": ["*"], + }, + } + + def my_tool(arg: str) -> str: + """A test tool.""" + return arg + + agent: GitHubCopilotAgent[GitHubCopilotOptions] = GitHubCopilotAgent( + client=mock_client, + tools=[my_tool], + default_options={ + "model": "gpt-5", + "provider": provider, + "mcp_servers": mcp_servers, + }, + ) + await agent.start() + + await agent._get_or_create_session(AgentSession()) # type: ignore + + call_args = mock_client.create_session.call_args + config = call_args.kwargs + assert config["provider"]["type"] == "azure" + assert config["model"] == "gpt-5" + assert config["mcp_servers"] is not None + assert config["tools"] is not None + + class TestGitHubCopilotAgentToolConversion: """Test cases for tool conversion.""" From ff05c22c5853a51b83c05b6fdb3b8e982bbf3b31 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Thu, 16 Apr 2026 00:23:37 +0200 Subject: [PATCH 4/9] Python: add experimental file history provider (#5248) * add experimental file history provider * Improve file history provider writes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * typo * cleanup * cleanup * fix in readme * added security messages * Refine file history provider locking Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * added additional sample --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/packages/core/AGENTS.md | 2 + .../packages/core/agent_framework/__init__.py | 2 + .../core/agent_framework/_feature_stage.py | 1 + .../core/agent_framework/_sessions.py | 265 +++++++++++++++++- .../packages/core/tests/core/test_sessions.py | 221 +++++++++++++++ .../samples/02-agents/conversations/README.md | 16 ++ .../conversations/file_history_provider.py | 157 +++++++++++ ...story_provider_conversation_persistence.py | 185 ++++++++++++ 8 files changed, 848 insertions(+), 1 deletion(-) create mode 100644 python/samples/02-agents/conversations/file_history_provider.py create mode 100644 python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index d6940289ac..30f946435a 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -63,6 +63,8 @@ agent_framework/ - **`SessionContext`** - Context object for session-scoped data during agent runs - **`ContextProvider`** - Base class for context providers (RAG, memory systems) - **`HistoryProvider`** - Base class for conversation history storage +- **`InMemoryHistoryProvider`** - Built-in session-state history provider for local runs +- **`FileHistoryProvider`** - JSON Lines file-backed history provider storing one file per session with one message record per line ### Skills (`_skills.py`) diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 497fc1496d..7475b1eb96 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -103,6 +103,7 @@ from ._middleware import ( from ._sessions import ( AgentSession, ContextProvider, + FileHistoryProvider, HistoryProvider, InMemoryHistoryProvider, SessionContext, @@ -318,6 +319,7 @@ __all__ = [ "FanInEdgeGroup", "FanOutEdgeGroup", "FileCheckpointStorage", + "FileHistoryProvider", "FinalT", "FinishReason", "FinishReasonLiteral", diff --git a/python/packages/core/agent_framework/_feature_stage.py b/python/packages/core/agent_framework/_feature_stage.py index 6fb698768c..1bda62b5d3 100644 --- a/python/packages/core/agent_framework/_feature_stage.py +++ b/python/packages/core/agent_framework/_feature_stage.py @@ -47,6 +47,7 @@ class ExperimentalFeature(str, Enum): """ EVALS = "EVALS" + FILE_HISTORY = "FILE_HISTORY" SKILLS = "SKILLS" diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 55d1a10a18..20125f19ff 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -8,16 +8,24 @@ This module provides the core types for the context provider pipeline: - HistoryProvider: Base class for history storage providers - AgentSession: Lightweight session state container - InMemoryHistoryProvider: Built-in in-memory history provider +- FileHistoryProvider: Built-in JSON Lines file history provider """ from __future__ import annotations +import asyncio import copy +import json +import threading import uuid +import weakref from abc import abstractmethod +from base64 import urlsafe_b64encode from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, ClassVar, TypeGuard, cast +from pathlib import Path +from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypeGuard, cast +from ._feature_stage import ExperimentalFeature, experimental from ._middleware import ChatContext, ChatMiddleware from ._types import AgentResponse, ChatResponse, Message, ResponseStream from .exceptions import ChatClientInvalidResponseException @@ -30,6 +38,17 @@ if TYPE_CHECKING: # Registry of known types for state deserialization _STATE_TYPE_REGISTRY: dict[str, type] = {} +JsonDumps: TypeAlias = Callable[[Any], str | bytes] +JsonLoads: TypeAlias = Callable[[str | bytes], Any] + + +def _default_json_dumps(value: Any) -> str: + return json.dumps(value, ensure_ascii=False) + + +def _default_json_loads(value: str | bytes) -> Any: + return json.loads(value) + def _is_middleware_sequence( middleware: MiddlewareTypes | Sequence[MiddlewareTypes], @@ -837,3 +856,247 @@ class InMemoryHistoryProvider(HistoryProvider): return existing = state.get("messages", []) state["messages"] = [*existing, *messages] + + +@experimental(feature_id=ExperimentalFeature.FILE_HISTORY) +class FileHistoryProvider(HistoryProvider): + """File-backed history provider that stores one JSON Lines file per session. + + Each persisted message is written as a single JSON object per line. The + provider does not serialize full session snapshots into the file. By default + it uses the standard library ``json`` module, but callers can inject + alternative ``dumps`` and ``loads`` callables compatible with the JSON + Lines format. + + Security posture: + Persisted history is stored as plaintext JSONL on the local filesystem. + Treat ``storage_path`` as trusted application storage, not as a secret + store. Encoded fallback filenames and resolved-path validation help + prevent path traversal via ``session_id``, but they do not encrypt file + contents or provide cross-process / cross-host locking. Use OS-level + file permissions, trusted directories, and carefully review what agent + or tool output is allowed to be persisted. + """ + + DEFAULT_SOURCE_ID: ClassVar[str] = "file_history" + DEFAULT_SESSION_FILE_STEM: ClassVar[str] = "default" + FILE_EXTENSION: ClassVar[str] = ".jsonl" + _FILE_LOCK_STRIPE_COUNT: ClassVar[int] = 64 + _ENCODED_SESSION_PREFIX: ClassVar[str] = "~session-" + _FILE_WRITE_LOCKS: ClassVar[tuple[threading.Lock, ...]] = tuple( + threading.Lock() for _ in range(_FILE_LOCK_STRIPE_COUNT) + ) + _WINDOWS_RESERVED_FILE_STEMS: ClassVar[frozenset[str]] = frozenset({ + "CON", + "PRN", + "AUX", + "NUL", + "COM1", + "COM2", + "COM3", + "COM4", + "COM5", + "COM6", + "COM7", + "COM8", + "COM9", + "LPT1", + "LPT2", + "LPT3", + "LPT4", + "LPT5", + "LPT6", + "LPT7", + "LPT8", + "LPT9", + }) + + def __init__( + self, + storage_path: str | Path, + *, + source_id: str = DEFAULT_SOURCE_ID, + load_messages: bool = True, + store_inputs: bool = True, + store_context_messages: bool = False, + store_context_from: set[str] | None = None, + store_outputs: bool = True, + skip_excluded: bool = False, + dumps: JsonDumps | None = None, + loads: JsonLoads | None = None, + ) -> None: + """Initialize the file history provider. + + Args: + storage_path: Directory path where session history files will be stored. + + Keyword Args: + source_id: Unique identifier for this provider instance. + load_messages: Whether to load messages before invocation. + store_inputs: Whether to store input messages. + store_context_messages: Whether to store context from other providers. + store_context_from: If set, only store context from these source_ids. + store_outputs: Whether to store response messages. + skip_excluded: When True, ``get_messages`` omits messages whose + ``additional_properties["_excluded"]`` is truthy. + dumps: Callable that serializes a message payload dict to JSON text + or UTF-8 bytes. The returned JSON must fit on a single line. + loads: Callable that deserializes JSON text or bytes back to a + message payload dict. + """ + super().__init__( + source_id=source_id, + load_messages=load_messages, + store_inputs=store_inputs, + store_context_messages=store_context_messages, + store_context_from=store_context_from, + store_outputs=store_outputs, + ) + self.storage_path = Path(storage_path) + self.storage_path.mkdir(parents=True, exist_ok=True) + self._storage_root = self.storage_path.resolve() + self.skip_excluded = skip_excluded + self.dumps = dumps or _default_json_dumps + self.loads = loads or _default_json_loads + self._async_write_locks_by_loop: weakref.WeakKeyDictionary[ + asyncio.AbstractEventLoop, + tuple[asyncio.Lock, ...], + ] = weakref.WeakKeyDictionary() + + async def get_messages( + self, + session_id: str | None, + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> list[Message]: + """Retrieve messages from the session's JSON Lines file.""" + del state, kwargs + file_path = self._session_file_path(session_id) + async_lock = self._session_async_write_lock(file_path) + thread_lock = self._session_write_lock(file_path) + + def _read_messages() -> list[Message]: + with thread_lock: + if not file_path.exists(): + return [] + + messages: list[Message] = [] + with file_path.open(encoding="utf-8") as file_handle: + for line_number, line in enumerate(file_handle, start=1): + serialized = line.strip() + if not serialized: + continue + try: + payload = self.loads(serialized) + except (TypeError, ValueError) as exc: + raise ValueError( + f"Failed to deserialize history line {line_number} from '{file_path}'." + ) from exc + if not isinstance(payload, Mapping): + raise ValueError( + f"History line {line_number} in '{file_path}' did not deserialize to a mapping." + ) + + try: + message = Message.from_dict(dict(cast(Mapping[str, Any], payload))) + except ValueError as exc: + raise ValueError( + f"History line {line_number} in '{file_path}' is not a valid Message payload." + ) from exc + messages.append(message) + return messages + + async with async_lock: + messages = await asyncio.to_thread(_read_messages) + if self.skip_excluded: + messages = [m for m in messages if not m.additional_properties.get("_excluded", False)] + return messages + + async def save_messages( + self, + session_id: str | None, + messages: Sequence[Message], + *, + state: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + """Append messages to the session's JSON Lines file.""" + del state, kwargs + if not messages: + return + + file_path = self._session_file_path(session_id) + async_lock = self._session_async_write_lock(file_path) + file_lock = self._session_write_lock(file_path) + + def _append_messages() -> None: + with file_lock, file_path.open("a", encoding="utf-8") as file_handle: + for message in messages: + file_handle.write(f"{self._serialize_message(message)}\n") + + async with async_lock: + await asyncio.to_thread(_append_messages) + + def _serialize_message(self, message: Message) -> str: + """Serialize a message payload to a single JSON Lines record.""" + serialized = self.dumps(message.to_dict()) + if isinstance(serialized, bytes): + serialized_text = serialized.decode("utf-8") + elif isinstance(serialized, str): + serialized_text = serialized + else: + raise TypeError("FileHistoryProvider.dumps must return str or bytes.") + + if "\n" in serialized_text or "\r" in serialized_text: + raise ValueError("FileHistoryProvider.dumps must return single-line JSON for JSON Lines storage.") + return serialized_text + + def _session_file_path(self, session_id: str | None) -> Path: + """Resolve the on-disk history file path for a session.""" + file_path = (self._storage_root / f"{self._session_file_stem(session_id)}{self.FILE_EXTENSION}").resolve() + if not file_path.is_relative_to(self._storage_root): + raise ValueError(f"Session history path escaped storage directory: {session_id!r}") + return file_path + + def _session_file_stem(self, session_id: str | None) -> str: + """Return the filename stem for a session.""" + raw_session_id = session_id or self.DEFAULT_SESSION_FILE_STEM + if self._is_literal_session_file_stem_safe(raw_session_id): + return raw_session_id + + encoded_session_id = urlsafe_b64encode(raw_session_id.encode("utf-8")).decode("ascii").rstrip("=") + return f"{self._ENCODED_SESSION_PREFIX}{encoded_session_id or self.DEFAULT_SESSION_FILE_STEM}" + + def _session_async_write_lock(self, file_path: Path) -> asyncio.Lock: + """Return the event-loop-local async lock for a session history file.""" + loop = asyncio.get_running_loop() + locks = self._async_write_locks_by_loop.get(loop) + if locks is None: + locks = tuple(asyncio.Lock() for _ in range(self._FILE_LOCK_STRIPE_COUNT)) + self._async_write_locks_by_loop[loop] = locks + return locks[self._lock_index(file_path)] + + @classmethod + def _session_write_lock(cls, file_path: Path) -> threading.Lock: + """Return the process-local thread lock for a session history file.""" + return cls._FILE_WRITE_LOCKS[cls._lock_index(file_path)] + + @classmethod + def _lock_index(cls, file_path: Path) -> int: + """Map a session history file to a bounded lock stripe.""" + return hash(file_path) % cls._FILE_LOCK_STRIPE_COUNT + + @classmethod + def _is_literal_session_file_stem_safe(cls, session_id: str) -> bool: + """Return whether the session ID can be used directly as a filename stem.""" + if ( + not session_id + or session_id.startswith(".") + or session_id.endswith((" ", ".")) + or session_id.upper() in cls._WINDOWS_RESERVED_FILE_STEMS + ): + return False + if any(ord(character) < 32 for character in session_id): + return False + return all(character.isalnum() or character in "._-" for character in session_id) diff --git a/python/packages/core/tests/core/test_sessions.py b/python/packages/core/tests/core/test_sessions.py index e5eacebfe5..ebb91d0b0d 100644 --- a/python/packages/core/tests/core/test_sessions.py +++ b/python/packages/core/tests/core/test_sessions.py @@ -1,7 +1,12 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import json +import threading +import time from collections.abc import Awaitable, Callable, Sequence +from pathlib import Path +from typing import Any import pytest @@ -10,6 +15,8 @@ from agent_framework import ( AgentSession, ChatContext, ContextProvider, + ExperimentalFeature, + FileHistoryProvider, HistoryProvider, InMemoryHistoryProvider, Message, @@ -505,3 +512,217 @@ class TestInMemoryHistoryProvider: ctx = SessionContext(session_id="s1", input_messages=[]) ctx.extend_messages("custom-source", [Message(role="user", contents=["test"])]) assert "custom-source" in ctx.context_messages + + +class TestFileHistoryProvider: + def test_is_marked_experimental(self) -> None: + assert FileHistoryProvider.__feature_stage__ == "experimental" + assert FileHistoryProvider.__feature_id__ == ExperimentalFeature.FILE_HISTORY.value + assert FileHistoryProvider.__doc__ is not None + assert ".. warning:: Experimental" in FileHistoryProvider.__doc__ + + async def test_stores_and_loads_messages(self, tmp_path: Path) -> None: + from agent_framework import AgentResponse + + provider = FileHistoryProvider(tmp_path) + session = AgentSession(session_id="s1") + + input_message = Message(role="user", contents=["hello"]) + response_message = Message(role="assistant", contents=["hi there"]) + first_context = SessionContext(session_id=session.session_id, input_messages=[input_message]) + + await provider.before_run( # type: ignore[arg-type] + agent=None, + session=session, + context=first_context, + state={}, + ) + first_context._response = AgentResponse(messages=[response_message]) + await provider.after_run( # type: ignore[arg-type] + agent=None, + session=session, + context=first_context, + state={}, + ) + + session_file = provider._session_file_path(session.session_id) + assert session_file.name == "s1.jsonl" + assert session_file.exists() + raw_lines = (await asyncio.to_thread(session_file.read_text, encoding="utf-8")).splitlines() + assert len(raw_lines) == 2 + payloads = [json.loads(line) for line in raw_lines] + assert all(payload["type"] == "message" for payload in payloads) + assert all("session_id" not in payload for payload in payloads) + + second_context = SessionContext( + session_id=session.session_id, input_messages=[Message(role="user", contents=["again"])] + ) + await provider.before_run( # type: ignore[arg-type] + agent=None, + session=session, + context=second_context, + state={}, + ) + loaded = second_context.context_messages.get(provider.source_id, []) + assert len(loaded) == 2 + assert loaded[0].text == "hello" + assert loaded[1].text == "hi there" + + def test_creates_storage_directory(self, tmp_path: Path) -> None: + nested_path = tmp_path / "nested" / "history" + provider = FileHistoryProvider(nested_path) + assert provider.storage_path == nested_path + assert nested_path.exists() + assert nested_path.is_dir() + + async def test_uses_encoded_filename_for_unsafe_session_id(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + unsafe_session_id = "../unsafe/session" + + await provider.save_messages(unsafe_session_id, [Message(role="user", contents=["hello"])]) + + session_file = provider._session_file_path(unsafe_session_id) + assert session_file.parent == provider.storage_path + assert session_file.name.startswith("~session-") + assert session_file.suffix == ".jsonl" + assert session_file.exists() + jsonl_files = await asyncio.to_thread( + lambda: sorted(path.name for path in provider.storage_path.glob("*.jsonl")) + ) + assert jsonl_files == [session_file.name] + + async def test_allows_custom_serializers_returning_bytes(self, tmp_path: Path) -> None: + calls: list[str] = [] + + def dumps(payload: object) -> bytes: + calls.append("dumps") + return json.dumps(payload).encode("utf-8") + + def loads(payload: str | bytes) -> object: + calls.append("loads") + if isinstance(payload, bytes): + payload = payload.decode("utf-8") + return json.loads(payload) + + provider = FileHistoryProvider(tmp_path, dumps=dumps, loads=loads) + + await provider.save_messages("custom-serializer", [Message(role="user", contents=["hello"])]) + loaded = await provider.get_messages("custom-serializer") + + assert calls == ["dumps", "loads"] + assert len(loaded) == 1 + assert loaded[0].text == "hello" + + async def test_invalid_jsonl_line_raises(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + await asyncio.to_thread(provider._session_file_path("broken").write_text, "{not-json}\n", encoding="utf-8") + + with pytest.raises(ValueError, match="Failed to deserialize history line 1"): + await provider.get_messages("broken") + + async def test_missing_session_file_returns_empty_messages(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + + loaded = await provider.get_messages("missing") + + assert loaded == [] + + async def test_none_session_id_uses_default_jsonl_file(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + + await provider.save_messages(None, [Message(role="user", contents=["hello"])]) + + session_file = provider._session_file_path(None) + assert session_file.name == "default.jsonl" + loaded = await provider.get_messages(None) + assert [message.text for message in loaded] == ["hello"] + + async def test_non_mapping_jsonl_line_raises(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path) + await asyncio.to_thread(provider._session_file_path("non-mapping").write_text, "[1, 2, 3]\n", encoding="utf-8") + + with pytest.raises(ValueError, match="did not deserialize to a mapping"): + await provider.get_messages("non-mapping") + + async def test_skip_excluded_omits_excluded_messages(self, tmp_path: Path) -> None: + provider = FileHistoryProvider(tmp_path, skip_excluded=True) + + await provider.save_messages( + "skip-excluded", + [ + Message(role="user", contents=["keep"]), + Message(role="assistant", contents=["skip"], additional_properties={"_excluded": True}), + ], + ) + + loaded = await provider.get_messages("skip-excluded") + + assert [message.text for message in loaded] == ["keep"] + + async def test_serializer_must_return_single_line_json(self, tmp_path: Path) -> None: + def dumps(payload: object) -> str: + return json.dumps(payload, indent=2) + + provider = FileHistoryProvider(tmp_path, dumps=dumps) + + with pytest.raises(ValueError, match="single-line JSON"): + await provider.save_messages("pretty-json", [Message(role="user", contents=["hello"])]) + + async def test_concurrent_writes_for_same_session_are_locked( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + provider = FileHistoryProvider(tmp_path) + session_id = "shared-session" + file_path = provider._session_file_path(session_id) + real_open = Path.open + write_started = threading.Event() + active_writes = 0 + overlap_detected = False + + class _TrackingFile: + def __init__(self, wrapped: Any) -> None: + self._wrapped = wrapped + + def __enter__(self) -> "_TrackingFile": + self._wrapped.__enter__() + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self._wrapped.__exit__(exc_type, exc_val, exc_tb) + + def write(self, data: str) -> int: + nonlocal active_writes, overlap_detected + write_started.set() + active_writes += 1 + overlap_detected = overlap_detected or active_writes > 1 + try: + time.sleep(0.05) + return int(self._wrapped.write(data)) + finally: + active_writes -= 1 + + def __getattr__(self, name: str) -> Any: + return getattr(self._wrapped, name) + + def tracked_open(path: Path, *args: Any, **kwargs: Any) -> Any: + handle = real_open(path, *args, **kwargs) + if path == file_path and args and args[0] == "a": + return _TrackingFile(handle) + return handle + + monkeypatch.setattr(Path, "open", tracked_open) + + first_save = asyncio.create_task(provider.save_messages(session_id, [Message(role="user", contents=["first"])])) + started = await asyncio.to_thread(write_started.wait, 1.0) + assert started + + second_save = asyncio.create_task( + provider.save_messages(session_id, [Message(role="assistant", contents=["second"])]) + ) + await asyncio.gather(first_save, second_save) + + assert not overlap_detected + loaded = await provider.get_messages(session_id) + assert [message.text for message in loaded] == ["first", "second"] diff --git a/python/samples/02-agents/conversations/README.md b/python/samples/02-agents/conversations/README.md index bbfb078659..002d4e6773 100644 --- a/python/samples/02-agents/conversations/README.md +++ b/python/samples/02-agents/conversations/README.md @@ -8,6 +8,8 @@ These samples demonstrate different approaches to managing conversation history |------|-------------| | [`suspend_resume_session.py`](suspend_resume_session.py) | Suspend and resume conversation sessions, comparing service-managed sessions (Azure AI Foundry) with in-memory sessions (OpenAI). | | [`custom_history_provider.py`](custom_history_provider.py) | Implement a custom history provider by extending `HistoryProvider`, enabling conversation persistence in your preferred storage backend. | +| [`file_history_provider.py`](file_history_provider.py) | Use the experimental `FileHistoryProvider` with `FoundryChatClient` and a function tool so the local JSON Lines file shows the full tool-calling loop. | +| [`file_history_provider_conversation_persistence.py`](file_history_provider_conversation_persistence.py) | Persist a tool-driven weather conversation with `FileHistoryProvider`, inspect the stored JSONL records, and continue with another city. | | [`cosmos_history_provider.py`](cosmos_history_provider.py) | Use Azure Cosmos DB as a history provider for durable conversation storage with `CosmosHistoryProvider`. | | [`cosmos_history_provider_conversation_persistence.py`](cosmos_history_provider_conversation_persistence.py) | Persist and resume conversations across application restarts using `CosmosHistoryProvider` — serialize session state, restore it, and continue with full Cosmos DB history. | | [`cosmos_history_provider_messages.py`](cosmos_history_provider_messages.py) | Direct message history operations — retrieve stored messages as a transcript, clear session history, and verify data deletion. | @@ -25,6 +27,20 @@ These samples demonstrate different approaches to managing conversation history **For `custom_history_provider.py`:** - `OPENAI_API_KEY`: Your OpenAI API key +**For `file_history_provider.py`:** +- `FOUNDRY_PROJECT_ENDPOINT`: Your Azure AI Foundry project endpoint +- `FOUNDRY_MODEL`: The Foundry model deployment name +- Azure CLI authentication (`az login`) +- The sample writes plaintext JSONL conversation logs to disk; use a trusted + local directory and avoid treating the history files as secure secret storage + +**For `file_history_provider_conversation_persistence.py`:** +- `FOUNDRY_PROJECT_ENDPOINT`: Your Azure AI Foundry project endpoint +- `FOUNDRY_MODEL`: The Foundry model deployment name +- Azure CLI authentication (`az login`) +- The sample writes plaintext JSONL conversation logs to disk; use a trusted + local directory and avoid treating the history files as secure secret storage + **For Cosmos DB samples (`cosmos_history_provider*.py`):** - `FOUNDRY_PROJECT_ENDPOINT`: Your Azure AI Foundry project endpoint - `FOUNDRY_MODEL`: The Foundry model deployment name diff --git a/python/samples/02-agents/conversations/file_history_provider.py b/python/samples/02-agents/conversations/file_history_provider.py new file mode 100644 index 0000000000..04a87f8224 --- /dev/null +++ b/python/samples/02-agents/conversations/file_history_provider.py @@ -0,0 +1,157 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import asyncio +import os +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Annotated + +# Uncomment this filter to suppress the experimental FileHistoryProvider warning +# before running the sample. +# import warnings # isort: skip +# warnings.filterwarnings("ignore", message=r"\[FILE_HISTORY\].*", category=FutureWarning) +from agent_framework import Agent, FileHistoryProvider, tool +from agent_framework.foundry import FoundryChatClient +from azure.identity import AzureCliCredential +from dotenv import load_dotenv +from pydantic import Field + +try: + import orjson +except ImportError: + orjson = None + + +# Load environment variables from .env file. +load_dotenv() + +""" +File History Provider + +This sample demonstrates how to use the experimental `FileHistoryProvider` with +`FoundryChatClient` and a function tool so the persisted JSON Lines file shows +the tool-calling loop as well as the regular chat turns. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT: Azure AI Foundry project endpoint. + FOUNDRY_MODEL: Foundry model deployment name. + +Key components: +- `FileHistoryProvider`: Stores one message JSON object per line in a local + `.jsonl` file for each session. +- `lookup_weather`: A function tool that makes the persisted file show the + assistant function call and tool result lines. +- `json.dumps(..., indent=2)`: Pretty-prints selected records in the sample + output while keeping the on-disk JSONL file compact and valid. +- `USE_TEMP_DIRECTORY`: Toggle between a temporary directory and a persistent + `sessions/` folder next to this sample file. + +Security posture: +- The history files are plaintext JSONL on disk, so use a trusted storage + directory and treat the files as conversation logs, not as secure secret + storage. +- Path safety checks protect the filename derived from the session id, but they + do not redact message contents or encrypt the file. +""" + +USE_TEMP_DIRECTORY = False +"""When True, store JSONL files in a temporary directory for this run only.""" + +LOCAL_SESSIONS_DIRECTORY_NAME = "sessions" +"""Folder name used when persisting history next to this sample file.""" + + +@tool(approval_mode="never_require") +def lookup_weather( + location: Annotated[str, Field(description="The city to look up weather for.")], +) -> str: + """Return a deterministic weather report for a city.""" + weather_reports = { + "Seattle": "Seattle is rainy with a high of 13C.", + "Amsterdam": "Amsterdam is cloudy with a high of 16C.", + } + return weather_reports.get(location, f"{location} is sunny with a high of 20C.") + + +@contextmanager +def _resolve_storage_directory() -> Iterator[Path]: + """Yield the configured storage directory for the sample run.""" + if USE_TEMP_DIRECTORY: + with tempfile.TemporaryDirectory(prefix="af-file-history-") as temp_directory: + yield Path(temp_directory) + return + + storage_directory = Path(__file__).resolve().parent / LOCAL_SESSIONS_DIRECTORY_NAME + storage_directory.mkdir(parents=True, exist_ok=True) + yield storage_directory + + +async def main() -> None: + """Run the file history provider sample.""" + + with _resolve_storage_directory() as storage_directory: + print(f"Using temporary directory: {USE_TEMP_DIRECTORY}") + print(f"Storage directory: {storage_directory}\n") + + # 2. Create the agent with a tool so the JSONL file includes tool-calling messages. + agent = Agent( + client=FoundryChatClient( + project_endpoint=os.getenv("FOUNDRY_PROJECT_ENDPOINT"), + model=os.getenv("FOUNDRY_MODEL"), + credential=AzureCliCredential(), + ), + name="FileHistoryAgent", + instructions=( + "You are a helpful assistant, use the lookup_weather tool for weather questions and " + "answer with the tool result in one sentence." + ), + tools=[lookup_weather], + # if orjson is available, use it for faster JSON serialization in the FileHistoryProvider, + # otherwise fall back to the default json module. + context_providers=[ + FileHistoryProvider( + storage_directory, + dumps=orjson.dumps if orjson else None, + loads=orjson.loads if orjson else None, + ) + ], + default_options={"store": False}, + ) + + # 3. Let Agent create the default UUID session id for this conversation. + session = agent.create_session() + + # 4. Ask a question that triggers the weather tool. + print("=== Run with tool calling ===") + query = "Use the lookup_weather tool for Seattle and tell me the weather." + response = await agent.run(query, session=session) + print(f"User: {query}") + print(f"Assistant: {response.text}\n") + + # 5. Ask a follow-up question that triggers the weather tool as well + print("=== Follow-up question ===") + query = "And what about Amsterdam?" + response = await agent.run(query, session=session) + print(f"User: {query}") + print(f"Assistant: {response.text}\n") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Sample output: +Using temporary directory: False +Storage directory: /path/to/samples/02-agents/conversations/sessions + +=== Run with tool calling === +User: Use the lookup_weather tool for Seattle and tell me the weather. +Assistant: +=== Follow-up question === +User: And what about Amsterdam? +Assistant: +""" diff --git a/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py b/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py new file mode 100644 index 0000000000..70c5d7e8e8 --- /dev/null +++ b/python/samples/02-agents/conversations/file_history_provider_conversation_persistence.py @@ -0,0 +1,185 @@ +# Copyright (c) Microsoft. All rights reserved. +# ruff: noqa: T201 + +from __future__ import annotations + +import asyncio +import json +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Annotated + +# Uncomment this filter to suppress the experimental FileHistoryProvider warning +# before running the sample. +# import warnings # isort: skip +# warnings.filterwarnings("ignore", message=r"\[FILE_HISTORY\].*", category=FutureWarning) +from agent_framework import Agent, FileHistoryProvider, tool +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from dotenv import load_dotenv +from pydantic import Field + +try: + import orjson +except ImportError: + orjson = None + + +load_dotenv() + +""" +File History Provider Conversation Persistence + +This sample demonstrates persisting a tool-driven conversation with the +experimental `FileHistoryProvider`, reading the stored JSONL file back from +disk, and then continuing the same conversation with another city. + +Environment variables: + FOUNDRY_PROJECT_ENDPOINT: Azure AI Foundry project endpoint. + FOUNDRY_MODEL: Foundry model deployment name. + +Key components: +- `FileHistoryProvider`: Stores one message JSON object per line in a local + `.jsonl` file for each session. +- `get_weather`: A function tool that makes the persisted file show the + assistant function call and tool result records. +- `json.dumps(..., indent=2)`: Pretty-prints a few persisted JSONL records + while keeping the on-disk file compact and valid. +- `load_dotenv()`: Loads `.env` values up front so the sample can stay focused + on history persistence instead of manual environment variable plumbing. +- Optional `orjson`: Uses `orjson.dumps` / `orjson.loads` automatically when + available, otherwise falls back to the standard library `json` module. + +Security posture: +- The history file is plaintext JSONL on disk, so use a trusted storage + directory and treat it as conversation logging, not as secure secret storage. +- Path safety checks protect the filename derived from the session id, but they + do not redact message contents or encrypt the file. +""" + +USE_TEMP_DIRECTORY = False +"""When True, store JSONL files in a temporary directory for this run only.""" + +LOCAL_SESSIONS_DIRECTORY_NAME = "sessions" +"""Folder name used when persisting history next to this sample file.""" + + +@tool(approval_mode="never_require") +def get_weather( + city: Annotated[str, Field(description="The city to get the weather for.")], +) -> str: + """Return a deterministic weather report for a city.""" + weather_reports = { + "Seattle": "Seattle is rainy with a high of 13C.", + "Amsterdam": "Amsterdam is cloudy with a high of 16C.", + } + return weather_reports.get(city, f"{city} is sunny with a high of 20C.") + + +@contextmanager +def _resolve_storage_directory() -> Iterator[Path]: + """Yield the configured storage directory for the sample run.""" + if USE_TEMP_DIRECTORY: + with tempfile.TemporaryDirectory(prefix="af-file-history-resume-") as temp_directory: + yield Path(temp_directory) + return + + storage_directory = Path(__file__).resolve().parent / LOCAL_SESSIONS_DIRECTORY_NAME + storage_directory.mkdir(parents=True, exist_ok=True) + yield storage_directory + + +async def main() -> None: + """Run the file history provider conversation persistence sample.""" + + with _resolve_storage_directory() as storage_directory: + print(f"Using temporary directory: {USE_TEMP_DIRECTORY}") + print(f"Storage directory: {storage_directory}\n") + + # 1. Create the client, history provider, and tool-enabled agent. + agent = Agent( + client=FoundryChatClient( + credential=AzureCliCredential(), + ), + name="WeatherHistoryAgent", + instructions=( + "You are a helpful assistant. Use the get_weather tool for weather questions " + "and answer in one sentence using the tool result." + ), + tools=[get_weather], + context_providers=[ + FileHistoryProvider( + storage_directory, + dumps=orjson.dumps if orjson else None, + loads=orjson.loads if orjson else None, + ) + ], + default_options={"store": False}, + ) + + # 2. Ask about the first city so the JSONL file is created on disk. + session = agent.create_session() + history_file = storage_directory / f"{session.session_id}.jsonl" + print("=== First weather question ===\n") + first_query = "Use the get_weather tool and tell me the weather in Seattle." + first_response = await agent.run(first_query, session=session) + print(f"User: {first_query}") + print(f"Assistant: {first_response.text}\n") + + # 3. Read the stored JSONL records back from disk and pretty-print a few of them. + raw_lines = (await asyncio.to_thread(history_file.read_text, encoding="utf-8")).splitlines() + print(f"Stored message lines after first question: {len(raw_lines)}") + print(f"History file: {history_file}\n") + print("=== JSONL preview from disk ===\n") + for index, line in enumerate(raw_lines[:4], start=1): + print(f"Record {index}:") + print(json.dumps(json.loads(line), indent=2)) + print() + + # 4. Continue the same persisted conversation with another city. + print("=== Second weather question ===\n") + second_query = "Now use the get_weather tool for Amsterdam." + second_response = await agent.run(second_query, session=session) + print(f"User: {second_query}") + print(f"Assistant: {second_response.text}\n") + + updated_lines = (await asyncio.to_thread(history_file.read_text, encoding="utf-8")).splitlines() + print(f"Stored message lines after second question: {len(updated_lines)}") + print(f"History file: {history_file}") + + +if __name__ == "__main__": + asyncio.run(main()) + +""" +Sample output: +Using temporary directory: False +Storage directory: /path/to/samples/02-agents/conversations/sessions + +=== First weather question === + +User: Use the get_weather tool and tell me the weather in Seattle. +Assistant: + +Stored message lines after first question: 4 +History file: /path/to/samples/02-agents/conversations/sessions/.jsonl + +=== JSONL preview from disk === + +Record 1: +{ + "type": "message", + "role": "user", + ... +} + +=== Second weather question === + +User: Now use the get_weather tool for Amsterdam. +Assistant: + +Stored message lines after second question: 8 +History file: /path/to/samples/02-agents/conversations/sessions/.jsonl +""" From f112150cfbc4d514b21b60a81bbe5239b4b2c81f Mon Sep 17 00:00:00 2001 From: Evan Mattson <35585003+moonbox3@users.noreply.github.com> Date: Thu, 16 Apr 2026 10:03:00 +0900 Subject: [PATCH 5/9] Python: bump misc-integration retry delay to 30s (#5293) The misc-integration job (Anthropic, Ollama, MCP) frequently fails on merge to main when the upstream MCP server (e.g. learn.microsoft.com/api/mcp) returns a transient rate-limit error. The previous 5s retry delay is too short to ride out the upstream backoff window, so all retries fail and the merge queue is blocked. Bumping to 30s gives the upstream a chance to recover before pytest-retry re-runs the test. --- .github/workflows/python-integration-tests.yml | 2 +- .github/workflows/python-merge-tests.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/python-integration-tests.yml b/.github/workflows/python-integration-tests.yml index 523a763b62..3c6c620614 100644 --- a/.github/workflows/python-integration-tests.yml +++ b/.github/workflows/python-integration-tests.yml @@ -171,7 +171,7 @@ jobs: -m integration -n logical --dist worksteal --timeout=120 --session-timeout=900 --timeout_method thread - --retries 2 --retry-delay 5 + --retries 2 --retry-delay 30 - name: Stop local MCP server if: always() shell: bash diff --git a/.github/workflows/python-merge-tests.yml b/.github/workflows/python-merge-tests.yml index 4fc47af595..454b297bed 100644 --- a/.github/workflows/python-merge-tests.yml +++ b/.github/workflows/python-merge-tests.yml @@ -287,7 +287,7 @@ jobs: -m integration -n logical --dist worksteal --timeout=120 --session-timeout=900 --timeout_method thread - --retries 2 --retry-delay 5 + --retries 2 --retry-delay 30 --junitxml=pytest.xml working-directory: ./python - name: Stop local MCP server From 611230cc8ebde031d6c15dbc15d7053ddf56b40c Mon Sep 17 00:00:00 2001 From: Evan Mattson <35585003+moonbox3@users.noreply.github.com> Date: Thu, 16 Apr 2026 11:34:28 +0900 Subject: [PATCH 6/9] Python: improve misc-integration test robustness (#5295) * Python: use local MCP server for hosted tools test and broaden image assertion The hosted tools integration test was hitting rate limits on the external learn.microsoft.com MCP server, causing persistent failures that retries couldn't recover from. Switch to the local MCP server already spun up in CI via LOCAL_MCP_URL, skipping when the env var isn't set. Also broaden the image description assertion to accept common synonyms (cottage, mansion, villa, etc.) instead of just "house", since the model legitimately uses varied vocabulary for the same image. * Address review feedback: validate LOCAL_MCP_URL scheme and use word boundaries - Skip hosted tools test when LOCAL_MCP_URL lacks http/https scheme, matching the pattern used in test_mcp.py. - Use regex word boundaries for image assertion to avoid false matches like "villain" matching "villa". --- .../anthropic/tests/test_anthropic_client.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index 52bb4c3a49..a10a8830b4 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. import os +import re from pathlib import Path from typing import Annotated, Any from unittest.mock import MagicMock, patch @@ -1503,6 +1504,10 @@ async def test_anthropic_client_integration_function_calling() -> None: @skip_if_anthropic_integration_tests_disabled async def test_anthropic_client_integration_hosted_tools() -> None: """Integration test for hosted tools.""" + local_mcp_url = os.environ.get("LOCAL_MCP_URL", "") + if not local_mcp_url or not local_mcp_url.startswith(("http://", "https://")): + pytest.skip("LOCAL_MCP_URL not set or not an HTTP URL; skipping hosted tools test") + client = AnthropicClient() messages = [Message(role="user", contents=["What tools do you have available?"])] @@ -1510,8 +1515,8 @@ async def test_anthropic_client_integration_hosted_tools() -> None: AnthropicClient.get_web_search_tool(), AnthropicClient.get_code_interpreter_tool(), AnthropicClient.get_mcp_tool( - name="example-mcp", - url="https://learn.microsoft.com/api/mcp", + name="local-mcp", + url=local_mcp_url, ), ] @@ -1607,7 +1612,8 @@ async def test_anthropic_client_integration_images() -> None: assert response is not None assert response.messages[0].text is not None - assert "house" in response.messages[0].text.lower() + text = response.messages[0].text.lower() + assert re.search(r"\b(house|home|building|cottage|mansion|villa)\b", text) # Response Format Tests From fe4cd3cddc99f157710296dad892bec427cae991 Mon Sep 17 00:00:00 2001 From: Evan Mattson <35585003+moonbox3@users.noreply.github.com> Date: Thu, 16 Apr 2026 11:46:49 +0900 Subject: [PATCH 7/9] Revert to public MCP server and skip on transient upstream errors (#5296) The local MCP server can't be used for hosted tools tests because Anthropic's backend needs to reach the MCP URL from their infrastructure (not localhost on the CI runner). Revert to learn.microsoft.com/api/mcp but catch BadRequestError, InternalServerError, APIConnectionError, and APITimeoutError and pytest.skip so upstream outages don't block the merge queue. --- .../anthropic/tests/test_anthropic_client.py | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index a10a8830b4..945e5356a4 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -1504,9 +1504,7 @@ async def test_anthropic_client_integration_function_calling() -> None: @skip_if_anthropic_integration_tests_disabled async def test_anthropic_client_integration_hosted_tools() -> None: """Integration test for hosted tools.""" - local_mcp_url = os.environ.get("LOCAL_MCP_URL", "") - if not local_mcp_url or not local_mcp_url.startswith(("http://", "https://")): - pytest.skip("LOCAL_MCP_URL not set or not an HTTP URL; skipping hosted tools test") + import anthropic client = AnthropicClient() @@ -1515,15 +1513,23 @@ async def test_anthropic_client_integration_hosted_tools() -> None: AnthropicClient.get_web_search_tool(), AnthropicClient.get_code_interpreter_tool(), AnthropicClient.get_mcp_tool( - name="local-mcp", - url=local_mcp_url, + name="example-mcp", + url="https://learn.microsoft.com/api/mcp", ), ] - response = await client.get_response( - messages=messages, - options={"tools": tools, "max_tokens": 100}, - ) + try: + response = await client.get_response( + messages=messages, + options={"tools": tools, "max_tokens": 100}, + ) + except ( + anthropic.BadRequestError, + anthropic.InternalServerError, + anthropic.APIConnectionError, + anthropic.APITimeoutError, + ) as e: + pytest.skip(f"Upstream MCP server unavailable: {e}") assert response is not None assert response.text is not None From 69697065ab78502c5e58a7e6bc90ae14fdc46c20 Mon Sep 17 00:00:00 2001 From: Chinedum Echeta <60179183+cecheta@users.noreply.github.com> Date: Thu, 16 Apr 2026 03:47:29 +0100 Subject: [PATCH 8/9] Python: Add context_providers and description to `workflow.as_agent()` (#4651) * Add context_providers and description to `workflow.as_agent()` * Add default workflow name and description * Positional * Move import --------- Co-authored-by: Tao Chen Co-authored-by: Evan Mattson <35585003+moonbox3@users.noreply.github.com> --- .../agent_framework/_workflows/_workflow.py | 29 ++++++++++++++--- .../tests/workflow/test_workflow_agent.py | 31 +++++++++++++++++++ 2 files changed, 55 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index 58050eece9..fc26db8953 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -11,11 +11,11 @@ import logging import types import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequence -from typing import Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, overload +from .._sessions import ContextProvider from .._types import ResponseStream from ..observability import OtelAttr, capture_exception, create_workflow_span -from ._agent import WorkflowAgent from ._checkpoint import CheckpointStorage from ._const import DEFAULT_MAX_ITERATIONS, GLOBAL_KWARGS_KEY, WORKFLOW_RUN_KWARGS_KEY from ._edge import ( @@ -35,6 +35,9 @@ from ._runner_context import RunnerContext from ._state import State from ._typing_utils import is_instance_of, try_coerce_to_type +if TYPE_CHECKING: + from ._agent import WorkflowAgent + logger = logging.getLogger(__name__) @@ -910,7 +913,14 @@ class Workflow(DictConvertible): return list(output_types) - def as_agent(self, name: str | None = None) -> WorkflowAgent: + def as_agent( + self, + name: str | None = None, + *, + description: str | None = None, + context_providers: Sequence[ContextProvider] | None = None, + **kwargs: Any, + ) -> WorkflowAgent: """Create a WorkflowAgent that wraps this workflow. The returned agent converts standard agent inputs (strings, Message, or lists of these) @@ -924,7 +934,10 @@ class Workflow(DictConvertible): initialization will fail with a ValueError. Args: - name: Optional name for the agent. If None, a default name will be generated. + name: Optional name for the agent. Defaults to workflow name. + description: Optional description of the agent. Defaults to workflow description. + context_providers: Optional sequence of context providers for the agent. + **kwargs: Additional keyword arguments passed to BaseAgent. Returns: A WorkflowAgent instance that wraps this workflow. @@ -935,4 +948,10 @@ class Workflow(DictConvertible): # Import here to avoid circular imports from ._agent import WorkflowAgent - return WorkflowAgent(workflow=self, name=name) + return WorkflowAgent( + workflow=self, + name=name if name is not None else self.name, + description=description if description is not None else self.description, + context_providers=context_providers, + **kwargs, + ) diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 9c7a655d23..0101a6e8a5 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -313,6 +313,37 @@ class TestWorkflowAgent: assert isinstance(agent_no_name, WorkflowAgent) assert agent_no_name.workflow is workflow + def test_workflow_as_agent_with_description_and_context_providers(self) -> None: + """Test that Workflow.as_agent() forwards description and context_providers.""" + executor = SimpleExecutor(id="executor1", response_text="Response") + workflow = WorkflowBuilder(start_executor=executor).build() + + history_provider = InMemoryHistoryProvider() + agent = workflow.as_agent( + name="MyAgent", + description="A test agent", + context_providers=[history_provider], + ) + + assert isinstance(agent, WorkflowAgent) + assert agent.name == "MyAgent" + assert agent.description == "A test agent" + assert history_provider in agent.context_providers + + def test_workflow_as_agent_defaults_name_and_description_from_workflow(self) -> None: + """Test that as_agent() defaults name and description to the workflow's own values.""" + executor = SimpleExecutor(id="executor1", response_text="Response") + workflow = WorkflowBuilder( + start_executor=executor, + name="my-workflow", + description="Workflow description", + ).build() + + agent = workflow.as_agent() + + assert agent.name == "my-workflow" + assert agent.description == "Workflow description" + def test_workflow_as_agent_cannot_handle_agent_inputs(self) -> None: """Test that Workflow.as_agent() raises an error if the start executor cannot handle agent inputs.""" From 8f7fd9525d1bf24f9606779ba7f8d41b66ce2ff1 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Wed, 15 Apr 2026 20:58:28 -0700 Subject: [PATCH 9/9] Python: Add OpenAI types to default checkpoint encoding allow list (#5297) * Add OpenAI types to default checkpoint encoding allow list * Address comments --- .../agent_framework/_workflows/_checkpoint.py | 8 ++-- .../_workflows/_checkpoint_encoding.py | 15 ++++-- .../test_checkpoint_unrestricted_pickle.py | 47 +++++++++++++++++++ 3 files changed, 61 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index f9a940a7db..22b4a1ea24 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -244,10 +244,10 @@ class FileCheckpointStorage: is serialized using pickle and embedded as base64-encoded strings within the JSON. This allows for human-readable checkpoint files while preserving the ability to store complex Python objects. - By default, checkpoint deserialization is restricted to a built-in set of safe - Python types (primitives, datetime, uuid, ...) and all ``agent_framework`` - internal types. To allow additional application-specific types, pass them via - the ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. + By default, checkpoint deserialization is restricted to a built-in set of safe Python types + (primitives, datetime, uuid, ...), all ``agent_framework`` internal types, and OpenAI SDK types + (``openai.types``). To allow additional application-specific types, pass them via the + ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. Example:: diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index a25a08c66a..dd1fb3d704 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -10,9 +10,9 @@ This hybrid approach provides: When ``allowed_types`` is supplied to :func:`decode_checkpoint_value`, a ``RestrictedUnpickler`` is used that limits which classes may be instantiated during deserialization. The default built-in safe set covers common Python -value types (primitives, datetime, uuid, ...) and all ``agent_framework`` -internal types. Callers can extend the set by passing additional -``"module:qualname"`` strings. +value types (primitives, datetime, uuid, ...), all ``agent_framework`` internal +types, and all ``openai.types`` types. Callers can extend the set by passing +additional ``"module:qualname"`` strings. """ from __future__ import annotations @@ -37,6 +37,9 @@ _JSON_NATIVE_TYPES = (str, int, float, bool, type(None)) # Module prefix for framework-internal types that are always allowed _FRAMEWORK_MODULE_PREFIX = "agent_framework." +# Module prefix for OpenAI SDK types that are always allowed +_OPENAI_MODULE_PREFIX = "openai.types." + # Built-in types considered safe for checkpoint deserialization. # Each entry is a ``module:qualname`` string matching the format produced by # :func:`_type_to_key`. These are the classes for which pickle's @@ -84,8 +87,9 @@ class _RestrictedUnpickler(pickle.Unpickler): # noqa: S301 """Unpickler that restricts which classes may be instantiated. Only classes whose ``module:qualname`` key appears in the combined allow - set (built-in safe types + framework types + caller-specified extras) are - permitted. All other classes raise :class:`pickle.UnpicklingError`. + set (built-in safe types + framework types + OpenAI SDK types + + caller-specified extras) are permitted. All other classes raise + :class:`pickle.UnpicklingError`. """ def __init__(self, data: bytes, allowed_types: frozenset[str]) -> None: @@ -99,6 +103,7 @@ class _RestrictedUnpickler(pickle.Unpickler): # noqa: S301 type_key in _BUILTIN_ALLOWED_TYPE_KEYS or type_key in self._allowed_types or module.startswith(_FRAMEWORK_MODULE_PREFIX) + or module.startswith(_OPENAI_MODULE_PREFIX) ): return super().find_class(module, name) # type: ignore[no-any-return] # nosec diff --git a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py index c70d8c85c3..77304841b2 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -216,3 +216,50 @@ def test_restricted_unpickler_raises_pickle_error(): unpickler = _RestrictedUnpickler(pickled, frozenset()) with pytest.raises(pickle.UnpicklingError, match="deserialization blocked"): unpickler.load() + + +def test_restricted_decode_allows_openai_types(): + """OpenAI SDK types are always allowed during restricted deserialization.""" + from openai.types.chat.chat_completion import ChatCompletion, Choice + from openai.types.chat.chat_completion_message import ChatCompletionMessage + from openai.types.completion_usage import CompletionUsage + + completion = ChatCompletion( + id="chatcmpl-test", + choices=[ + Choice( + finish_reason="stop", + index=0, + message=ChatCompletionMessage(role="assistant", content="hello"), + ) + ], + created=1700000000, + model="gpt-4", + object="chat.completion", + usage=CompletionUsage(completion_tokens=1, prompt_tokens=1, total_tokens=2), + ) + encoded = encode_checkpoint_value(completion) + decoded = decode_checkpoint_value(encoded, allowed_types=frozenset()) + + assert isinstance(decoded, ChatCompletion) + assert decoded.id == "chatcmpl-test" + assert decoded.choices[0].message.content == "hello" + + +def test_restricted_decode_allows_openai_response_types(): + """OpenAI Responses API types are always allowed during restricted deserialization.""" + from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails, ResponseUsage + + usage = ResponseUsage( + input_tokens=10, + output_tokens=20, + total_tokens=30, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + ) + encoded = encode_checkpoint_value(usage) + decoded = decode_checkpoint_value(encoded, allowed_types=frozenset()) + + assert isinstance(decoded, ResponseUsage) + assert decoded.input_tokens == 10 + assert decoded.output_tokens == 20