Python: [BREAKING] Moved to a single get_response and run API (#3379)

* WIP * big update to new ResponseStream model * fixed tests and typing * fixed tests and typing * fixed tools typevar import * fix * mypy fix * mypy fixes and some cleanup * fix missing quoted names * and client * fix imports agui * fix anthropic override * fix agui * fix ag ui * fix import * fix anthropic types * fix mypy * refactoring * updated typing * fix 3.11 * fixes * redid layering of chat clients and agents * redid layering of chat clients and agents * Fix lint, type, and test issues after rebase - Add @overload decorators to AgentProtocol.run() for type compatibility - Add missing docstring params (middleware, function_invocation_configuration) - Fix TODO format (TD002) by adding author tags - Fix broken observability tests from upstream: - Replace non-existent use_instrumentation with direct instantiation - Replace non-existent use_agent_instrumentation with AgentTelemetryLayer mixin - Fix get_streaming_response to use get_response(stream=True) - Add AgentInitializationError import - Update streaming exception tests to match actual behavior * Fix AgentExecutionException import error in test_agents.py - Replace non-existent AgentExecutionException with AgentRunException * Fix test import and asyncio deprecation issues - Add 'tests' to pythonpath in ag-ui pyproject.toml for utils_test_ag_ui import - Replace deprecated asyncio.get_event_loop().run_until_complete with asyncio.run * Fix azure-ai test failures - Update _prepare_options patching to use correct class path - Fix test_to_azure_ai_agent_tools_web_search_missing_connection to clear env vars * Convert ag-ui utils_test_ag_ui.py to conftest.py - Move test utilities to conftest.py for proper pytest discovery - Update all test imports to use conftest instead of utils_test_ag_ui - Remove old utils_test_ag_ui.py file - Revert pythonpath change in pyproject.toml * fix: use relative imports for ag-ui test utilities * fix agui * Rename Bare*Client to Raw*Client and BaseChatClient - Renamed BareChatClient to BaseChatClient (abstract base class) - Renamed BareOpenAIChatClient to RawOpenAIChatClient - Renamed BareOpenAIResponsesClient to RawOpenAIResponsesClient - Renamed BareAzureAIClient to RawAzureAIClient - Added warning docstrings to Raw* classes about layer ordering - Updated README in samples/getting_started/agents/custom with layer docs - Added test for span ordering with function calling * Fix layer ordering: FunctionInvocationLayer before ChatTelemetryLayer This ensures each inner LLM call gets its own telemetry span, resulting in the correct span sequence: chat -> execute_tool -> chat Updated all production clients and test mocks to use correct ordering: - ChatMiddlewareLayer (first) - FunctionInvocationLayer (second) - ChatTelemetryLayer (third) - BaseChatClient/Raw...Client (fourth) * Remove run_stream usage * Fix conversation_id propagation * Python: Add BaseAgent implementation for Claude Agent SDK (#3509) * Added ClaudeAgent implementation * Updated streaming logic * Small updates * Small update * Fixes * Small fix * Naming improvements * Updated imports * Addressed comments * Updated package versions * Update Claude agent connector layering * fix test and plugin * Store function middleware in invocation layer * Fix telemetry streaming and ag-ui tests * Remove legacy ag-ui tests folder * updates * Remove terminate flag from FunctionInvocationContext, use MiddlewareTermination instead - Remove terminate attribute from FunctionInvocationContext - Add result attribute to MiddlewareTermination to carry function results - FunctionMiddlewarePipeline.execute() now lets MiddlewareTermination propagate - _auto_invoke_function captures context.result in exception before re-raising - _try_execute_function_calls catches MiddlewareTermination and sets should_terminate - Fix handoff middleware to append to chat_client.function_middleware directly - Update tests to use raise MiddlewareTermination instead of context.terminate - Add middleware flow documentation in samples/concepts/tools/README.md - Fix ag-ui to use FunctionMiddlewarePipeline instead of removed create_function_middleware_pipeline * fix: remove references to removed terminate flag in purview tests, add type ignore * fix: move _test_utils.py from package to test folder * fix: call get_final_response() to trigger context provider notification in streaming test * fix: correct broken links in tools README * docs: clarify default middleware behavior in summary table * fix: ensure inner stream result hooks are called when using map()/from_awaitable() * Fix mypy type errors * Address PR review comments on observability.py - Remove TODO comment about unconsumed streams, add explanatory note instead - Remove redundant _close_span cleanup hook (already called in _finalize_stream) - Clarify behavior: cleanup hooks run after stream iteration, if stream is not consumed the span remains open until garbage collected * Remove gen_ai.client.operation.duration from span attributes Duration is a metrics-only attribute per OpenTelemetry semantic conventions. It should be recorded to the histogram but not set as a span attribute. * Remove duration from _get_response_attributes, pass directly to _capture_response Duration is a metrics-only attribute. It's now passed directly to _capture_response instead of being included in the attributes dict that gets set on the span. * Remove redundant _close_span cleanup hook in AgentTelemetryLayer _finalize_stream already calls _close_span() in its finally block, so adding it as a separate cleanup hook is redundant. * Use weakref.finalize to close span when stream is garbage collected If a user creates a streaming response but never consumes it, the cleanup hooks won't run. Now we register a weak reference finalizer that will close the span when the stream object is garbage collected, ensuring spans don't leak in this scenario. * Fix _get_finalizers_from_stream to use _result_hooks attribute Renamed function to _get_result_hooks_from_stream and fixed it to look for the _result_hooks attribute which is the correct name in ResponseStream class. * Add missing asyncio import in test_request_info_mixin.py * Fix leftover merge conflict marker in image_generation sample * Update integration tests * Fix integration tests: increase max_iterations from 1 to 2 Tests with tool_choice options require at least 2 iterations: 1. First iteration to get function call and execute the tool 2. Second iteration to get the final text response With max_iterations=1, streaming tests would return early with only the function call/result but no final text content. * Fix duplicate function call error in conversation-based APIs When using conversation_id (for Responses/Assistants APIs), the server already has the function call message from the previous response. We should only send the new function result message, not all messages including the function call which would cause a duplicate ID error. Fix: When conversation_id is set, only send the last message (the tool result) instead of all response.messages. * Add regression test for conversation_id propagation between tool iterations Port test from PR #3664 with updates for new streaming API pattern. Tests that conversation_id is properly updated in options dict during function invocation loop iterations. * Fix tool_choice=required to return after tool execution When tool_choice is 'required', the user's intent is to force exactly one tool call. After the tool executes, return immediately with the function call and result - don't continue to call the model again. This fixes integration tests that were failing with empty text responses because with tool_choice=required, the model would keep returning function calls instead of text. Also adds regression tests for: - conversation_id propagation between tool iterations (from PR #3664) - tool_choice=required returns after tool execution * Document tool_choice behavior in tools README - Add table explaining tool_choice values (auto, none, required) - Explain why tool_choice=required returns immediately after tool execution - Add code example showing the difference between required and auto - Update flow diagram to show the early return path for tool_choice=required * Fix tool_choice=None behavior - don't default to 'auto' Remove the hardcoded default of 'auto' for tool_choice in ChatAgent init. When tool_choice is not specified (None), it will now not be sent to the API, allowing the API's default behavior to be used. Users who want tool_choice='auto' can still explicitly set it either in default_options or at runtime. Fixes #3585 * Fix tool_choice=none should not remove tools In OpenAI Assistants client, tools were not being sent when tool_choice='none'. This was incorrect - tool_choice='none' means the model won't call tools, but tools should still be available in the request (they may be used later in the conversation). Fixes #3585 * Add test for tool_choice=none preserving tools Adds a regression test to ensure that when tool_choice='none' is set but tools are provided, the tools are still sent to the API. This verifies the fix for #3585. * Fix tool_choice=none should not remove tools in all clients Apply the same fix to OpenAI Responses client and Azure AI client: - OpenAI Responses: Remove else block that popped tool_choice/parallel_tool_calls - Azure AI: Remove tool_choice != 'none' check when adding tools When tool_choice='none', the model won't call tools, but tools should still be sent to the API so they're available for future turns. Also update README to clarify tool_choice=required supports multiple tools. Fixes #3585 * Keep tool_choice even when tools is None Move tool_choice processing outside of the 'if tools' block in OpenAI Responses client so tool_choice is sent to the API even when no tools are provided. * Update test to match new parallel_tool_calls behavior Changed test_prepare_options_removes_parallel_tool_calls_when_no_tools to test_prepare_options_preserves_parallel_tool_calls_when_no_tools to reflect that parallel_tool_calls is now preserved even when no tools are present, consistent with the tool_choice behavior. * Fix ChatMessage API and Role enum usage after rebase - Update ChatMessage instantiation to use keyword args (role=, text=, contents=) - Fix Role enum comparisons to use .value for string comparison - Add created_at to AgentResponse in error handling - Fix AgentResponse.from_updates -> from_agent_run_response_updates - Fix DurableAgentStateMessage.from_chat_message to convert Role enum to string - Add Role import where needed * Fix additional ChatMessage API and method name changes - Fix ChatMessage usage in workflow files (use text= instead of contents= for strings) - Fix AgentResponse.from_updates -> from_agent_run_response_updates in workflow files - Fix test files for ChatMessage and Role enum usage * Fix remaining ChatMessage API usage in test files * Fix more ChatMessage and Role API changes in source and test files - Fix ChatMessage in _magentic.py replan method - Fix Role enum comparison in test assertions - Fix remaining test files with old ChatMessage syntax * Fix ChatMessage and Role API changes across packages - Add Role import where missing - Fix ChatMessage signature: positional args to keyword args (role=, text=, contents=) - Fix Role enum comparisons: .role.value instead of .role string - Fix FinishReason enum usage in ag-ui event converters - Rename AgentResponse.from_updates to from_agent_run_response_updates in ag-ui Fixes API compatibility after Types API Review improvements merge * Fix ChatMessage and Role API changes in github_copilot tests * Fix ChatMessage and Role API changes in redis and github_copilot packages - Fix redis provider: Role enum comparison using .value - Fix redis tests: ChatMessage signature and Role comparisons - Fix github_copilot tests: ChatMessage signature and Role comparisons - Update docstring examples in redis chat message store * Fix ChatMessage and Role API changes in devui package - Fix executor: ChatMessage signature change - Fix conversations: Role enum to string conversion in two places - Fix tests: ChatMessage signatures and Role comparisons * Fix ChatMessage and Role API changes in a2a and lab packages - Fix a2a tests: Role comparisons and ChatMessage signatures - Fix lab tau2 source: Role enum comparison in flip_messages, log_messages, sliding_window - Fix lab tau2 tests: ChatMessage signatures and Role comparisons * Remove duplicate test files from ag-ui/tests (tests are in ag_ui_tests) * Fix ChatMessage and Role API changes across packages After rebasing on upstream/main which merged PR #3647 (Types API Review improvements), fix all packages to use the new API: - ChatMessage: Use keyword args (role=, text=, contents=) instead of positional args - Role: Compare using .value attribute since it's now an enum Packages fixed: - ag-ui: Fixed Role value extraction bugs in _message_adapters.py - anthropic: Fixed ChatMessage and Role comparisons in tests - azure-ai: Fixed Role comparison in _client.py - azure-ai-search: Fixed ChatMessage and Role in source/tests - bedrock: Fixed ChatMessage signatures in tests - chatkit: Fixed ChatMessage and Role in source/tests - copilotstudio: Fixed ChatMessage and Role in tests - declarative: Fixed ChatMessage in _executors_agents.py - mem0: Fixed ChatMessage and Role in source/tests - purview: Fixed ChatMessage in source/tests * Fix mypy errors for ChatMessage and Role API changes - durabletask: Use str() fallback in role value extraction - core: Fix ChatMessage in _orchestrator_helpers.py to use keyword args - core: Add type ignore for _conversation_state.py contents deserialization - ag-ui: Fix type ignore comments (call-overload instead of arg-type) - azure-ai-search: Fix get_role_value type hint to accept Any - lab: Move get_role_value to module level with Any type hint * Improve CI test timeout configuration - Increase job timeout from 10 to 15 minutes - Reduce per-test timeout to 60s (was 900s/300s) - Add --timeout_method thread for better timeout handling - Add --timeout-verbose to see which tests are slow - Reduce retries from 3 to 2 and delay from 10s to 5s This ensures individual test timeouts are shorter than the job timeout, providing better visibility when tests hang. With 60s timeout and 2 retries, worst case per test is ~180s. * Fix ChatMessage API usage in docstrings and source - Fix ChatMessage positional args in docstrings: _serialization.py, _threads.py, _middleware.py - Fix ChatMessage in tau2 runner.py - Fix role comparison in _orchestrator_helpers.py to use .value - Fix role comparison in _group_chat.py docstring example - Fix role assertions in test_durable_entities.py to use .value * Revert tool_choice/parallel_tool_calls changes - must be removed when no tools OpenAI API requires tool_choice and parallel_tool_calls to only be present when tools are specified. Restored the logic that removes these options when there are no tools. - Restored check in _chat_client.py to remove tool_choice and parallel_tool_calls when no tools present - Restored same logic in _responses_client.py - Reverted test to expect the correct behavior * fixed issue in tests * fix: resolve merge conflict markers in ag-ui tests * fix: restructure ag-ui tests and fix Role/FinishReason to use string types * fix: streaming function invocation and middleware termination - Refactor streaming function invocation to use get_final_response() on inner streams - Fix MiddlewareTermination to accept result parameter for passing results - Fix _AutoHandoffMiddleware to use MiddlewareTermination instead of context.terminate - Fix AgentMiddlewareLayer.run() to properly forward function/chat middleware - Remove duplicate middleware registration in AgentMiddlewareLayer.__init__ - Fix exception handling in _auto_invoke_function to properly capture termination - Fix mypy errors in core package - Update tests to use stream=True parameter for unified run API * fix all tests command * Refactor integration tests to use pytest fixtures - Merge testutils.py into conftest.py for azurefunctions integration tests - Merge dt_testutils.py into conftest.py for durabletask integration tests - Convert all integration tests to use fixtures instead of direct imports (fixes ModuleNotFoundError with --import-mode=importlib) - Add sample_helper fixture for azurefunctions tests - Add agent_client_factory and orchestration_helper fixtures for durabletask - Integration tests now skip with descriptive messages when services unavailable - Restructure devui tests into tests/devui/ with proper conftest.py - Add test organization guidelines to CODING_STANDARD.md - Remove __init__.py from test directories per pytest best practices * Fix pytest_collection_modifyitems to only skip integration tests The hook was skipping all tests in the test session, not just integration tests. Now it only skips items in the integration_tests directory. * Fix mem0 tests failing on Python 3.13 Use patch.object on the imported module instead of @patch with string path to ensure the mock takes effect regardless of import timing. * fix mem0 * another attempt for mem0 * fix for mem0 * fix mem0 * Increase worker initialization wait time in durabletask tests Increase from 2 to 8 seconds to allow time for: - Python startup and module imports - Azure OpenAI client creation - Agent registration with DTS worker - Worker connection to DTS This helps prevent test failures in CI where the first tests may run before the worker is fully ready to process requests. * Fix streaming test to use ResponseStream with finalizer The _consume_stream method now expects a ResponseStream that can provide a final AgentResponse via get_final_response(). Update the test to use ResponseStream with AgentResponse.from_updates as the finalizer. * Fix MockToolCallingAgent to use new ResponseStream API and update samples * small updates to run_stream to run * fix sub workflow * temp fix for az func test --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-02-05 21:09:58 +01:00
parent d1205896a1
commit 3dc59c83b5
372 changed files with 11583 additions and 9465 deletions
@@ -303,7 +303,7 @@ class InMemoryConversationStore(ConversationStore):
            content = item.get("content", [])
            text = content[0].get("text", "") if content else ""

-            chat_msg = ChatMessage(role, [{"type": "text", "text": text}])
+            chat_msg = ChatMessage(role=role, text=text)  # type: ignore[arg-type]
            chat_messages.append(chat_msg)

        # Add messages to AgentThread
@@ -588,7 +588,7 @@ class InMemoryConversationStore(ConversationStore):
        return None

    def get_thread(self, conversation_id: str) -> AgentThread | None:
-        """Get AgentThread for execution - CRITICAL for agent.run_stream()."""
+        """Get AgentThread for execution - CRITICAL for agent.run()."""
        conv_data = self._conversations.get(conversation_id)
        return conv_data["thread"] if conv_data else None

@@ -111,7 +111,7 @@ class EntityDiscovery:
                f"Only 'directory' and 'in-memory' sources are supported."
            )

-        # Note: Checkpoint storage is now injected at runtime via run_stream() parameter,
+        # Note: Checkpoint storage is now injected at runtime via run() parameter,
        # not at load time. This provides cleaner architecture and explicit control flow.
        # See _executor.py _execute_workflow() for runtime checkpoint storage injection.

@@ -361,16 +361,10 @@ class EntityDiscovery:

        # Log helpful info about agent capabilities (before creating EntityInfo)
        if entity_type == "agent":
-            has_run_stream = hasattr(entity_object, "run_stream")
            has_run = hasattr(entity_object, "run")

-            if not has_run_stream and has_run:
-                logger.info(
-                    f"Agent '{entity_id}' only has run() (non-streaming). "
-                    "DevUI will automatically convert to streaming."
-                )
-            elif not has_run_stream and not has_run:
-                logger.warning(f"Agent '{entity_id}' lacks both run() and run_stream() methods. May not work.")
+            if not has_run:
+                logger.warning(f"Agent '{entity_id}' lacks run() method. May not work.")

        # Check deployment support based on source
        # For directory-based entities, we need the path to verify deployment support
@@ -407,7 +401,6 @@ class EntityDiscovery:
                "class_name": entity_object.__class__.__name__
                if hasattr(entity_object, "__class__")
                else str(type(entity_object)),
-                "has_run_stream": hasattr(entity_object, "run_stream"),
            },
        )

@@ -774,9 +767,9 @@ class EntityDiscovery:
                pass

            # Fallback to duck typing for agent protocol
-            # Agent must have either run_stream() or run() method, plus id and name
-            has_execution_method = hasattr(obj, "run_stream") or hasattr(obj, "run")
-            if has_execution_method and hasattr(obj, "id") and hasattr(obj, "name"):
+            # Agent must have run() method, plus id and name
+            has_run = hasattr(obj, "run")
+            if has_run and hasattr(obj, "id") and hasattr(obj, "name"):
                return True

        except (TypeError, AttributeError):
@@ -793,8 +786,9 @@ class EntityDiscovery:
        Returns:
            True if object appears to be a valid workflow
        """
-        # Check for workflow - must have run_stream method and executors
-        return hasattr(obj, "run_stream") and (hasattr(obj, "executors") or hasattr(obj, "get_executors_list"))
+        # Check for workflow - must have run (streaming via stream=True) and executors
+        has_run = hasattr(obj, "run")
+        return has_run and (hasattr(obj, "executors") or hasattr(obj, "get_executors_list"))

    async def _register_entity_from_object(
        self, obj: Any, obj_type: str, module_path: str, source: str = "directory"
@@ -858,7 +852,6 @@ class EntityDiscovery:
                    "module_path": module_path,
                    "entity_type": obj_type,
                    "source": source,
-                    "has_run_stream": hasattr(obj, "run_stream"),
                    "class_name": obj.__class__.__name__ if hasattr(obj, "__class__") else str(type(obj)),
                },
            )
@@ -326,37 +326,23 @@ class AgentFrameworkExecutor:
            # but is_connected stays True. Detect and reconnect before execution.
            await self._ensure_mcp_connections(agent)

-            # Check if agent supports streaming
-            if hasattr(agent, "run_stream") and callable(agent.run_stream):
-                # Use Agent Framework's native streaming with optional thread
+            # Agent must have run() method - use stream=True for streaming
+            if hasattr(agent, "run") and callable(agent.run):
+                # Use Agent Framework's run() with stream=True for streaming
                if thread:
-                    async for update in agent.run_stream(user_message, thread=thread):
+                    async for update in agent.run(user_message, stream=True, thread=thread):
                        for trace_event in trace_collector.get_pending_events():
                            yield trace_event

                        yield update
                else:
-                    async for update in agent.run_stream(user_message):
+                    async for update in agent.run(user_message, stream=True):
                        for trace_event in trace_collector.get_pending_events():
                            yield trace_event

                        yield update
-            elif hasattr(agent, "run") and callable(agent.run):
-                # Non-streaming agent - use run() and yield complete response
-                logger.info("Agent lacks run_stream(), using run() method (non-streaming)")
-                if thread:
-                    response = await agent.run(user_message, thread=thread)
-                else:
-                    response = await agent.run(user_message)
-
-                # Yield trace events before response
-                for trace_event in trace_collector.get_pending_events():
-                    yield trace_event
-
-                # Yield the complete response (mapper will convert to streaming events)
-                yield response
            else:
-                raise ValueError("Agent must implement either run() or run_stream() method")
+                raise ValueError("Agent must implement run() method")

            # Emit agent lifecycle completion event
            from .models._openai_custom import AgentCompletedEvent
@@ -426,7 +412,7 @@ class AgentFrameworkExecutor:

            # Get session-scoped checkpoint storage (InMemoryCheckpointStorage from conv_data)
            # Each conversation has its own storage instance, providing automatic session isolation.
-            # This storage is passed to workflow.run_stream() which sets it as runtime override,
+            # This storage is passed to workflow.run(stream=True) which sets it as runtime override,
            # ensuring all checkpoint operations (save/load) use THIS conversation's storage.
            # The framework guarantees runtime storage takes precedence over build-time storage.
            checkpoint_storage = self.checkpoint_manager.get_checkpoint_storage(conversation_id)
@@ -478,15 +464,17 @@ class AgentFrameworkExecutor:
                # NOTE: Two-step approach for stateless HTTP (framework limitation):
                # 1. Restore checkpoint to load pending requests into workflow's in-memory state
                # 2. Then send responses using send_responses_streaming
-                # Future: Framework should support run_stream(checkpoint_id, responses) in single call
+                # Future: Framework should support run(stream=True, checkpoint_id, responses) in single call
                # (checkpoint_id is guaranteed to exist due to earlier validation)
                logger.debug(f"Restoring checkpoint {checkpoint_id} then sending HIL responses")

                try:
                    # Step 1: Restore checkpoint to populate workflow's in-memory pending requests
                    restored = False
-                    async for _event in workflow.run_stream(
-                        checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage
+                    async for _event in workflow.run(
+                        stream=True,
+                        checkpoint_id=checkpoint_id,
+                        checkpoint_storage=checkpoint_storage,
                    ):
                        restored = True
                        break  # Stop immediately after restoration, don't process events
@@ -545,8 +533,10 @@ class AgentFrameworkExecutor:
                logger.info(f"Resuming workflow from checkpoint {checkpoint_id} in session {conversation_id}")

                try:
-                    async for event in workflow.run_stream(
-                        checkpoint_id=checkpoint_id, checkpoint_storage=checkpoint_storage
+                    async for event in workflow.run(
+                        stream=True,
+                        checkpoint_id=checkpoint_id,
+                        checkpoint_storage=checkpoint_storage,
                    ):
                        if isinstance(event, RequestInfoEvent):
                            self._enrich_request_info_event_with_response_schema(event, workflow)
@@ -571,7 +561,7 @@ class AgentFrameworkExecutor:

                parsed_input = await self._parse_workflow_input(workflow, request.input)

-                async for event in workflow.run_stream(parsed_input, checkpoint_storage=checkpoint_storage):
+                async for event in workflow.run(parsed_input, stream=True, checkpoint_storage=checkpoint_storage):
                    if isinstance(event, RequestInfoEvent):
                        self._enrich_request_info_event_with_response_schema(event, workflow)

@@ -760,7 +750,7 @@ class AgentFrameworkExecutor:
        if not contents:
            contents.append(Content.from_text(text=""))

-        chat_message = ChatMessage("user", contents)
+        chat_message = ChatMessage(role="user", contents=contents)

        logger.info(f"Created ChatMessage with {len(contents)} contents:")
        for idx, content in enumerate(contents):
@@ -161,7 +161,7 @@ export function AgentDetailsModal({
            </DetailCard>
          )}

-          {/* Tools and Middleware Grid */}
+          {/* Tools and MiddlewareTypes Grid */}
          <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
            {/* Tools */}
            {agent.tools && agent.tools.length > 0 && (
@@ -30,7 +30,7 @@ dependencies = [
 ]

 [project.optional-dependencies]
-dev = ["pytest>=7.0.0", "watchdog>=3.0.0"]
+dev = ["pytest>=7.0.0", "watchdog>=3.0.0", "agent-framework-orchestrations"]
 all = ["pytest>=7.0.0", "watchdog>=3.0.0"]

 [project.scripts]
@@ -49,7 +49,7 @@ fallback-version = "0.0.0"

 [tool.pytest.ini_options]
 testpaths = 'tests'
-pythonpath = ["tests"]
+pythonpath = ["tests/devui"]
 addopts = "-ra -q -r fEX"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
@@ -1,22 +1,21 @@
 # Copyright (c) Microsoft. All rights reserved.

-"""Shared test utilities for DevUI tests.
+"""Pytest configuration and fixtures for DevUI tests.

-This module provides reusable test helpers including:
+This module provides reusable test fixtures including:
 - Mock chat clients that don't require API keys
 - Real workflow event classes from agent_framework
 - Test agents and executors for workflow testing
 - Factory functions for test data
-
-These follow the patterns established in other agent_framework packages
-(like a2a, ag-ui) which use explicit imports instead of conftest.py
-to avoid pytest plugin conflicts when running tests across packages.
 """

 import sys
-from collections.abc import AsyncIterable, MutableSequence
+from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
+from pathlib import Path
 from typing import Any, Generic

+import pytest
+import pytest_asyncio
 from agent_framework import (
    AgentResponse,
    AgentResponseUpdate,
@@ -28,30 +27,29 @@ from agent_framework import (
    ChatResponse,
    ChatResponseUpdate,
    Content,
-    use_chat_middleware,
+    ResponseStream,
 )
 from agent_framework._clients import TOptions_co
 from agent_framework._workflows._agent_executor import AgentExecutorResponse
-from agent_framework.orchestrations import ConcurrentBuilder, SequentialBuilder
-
-if sys.version_info >= (3, 12):
-    from typing import override  # type: ignore # pragma: no cover
-else:
-    from typing_extensions import override  # type: ignore[import] # pragma: no cover
-
-# Import real workflow event classes - NOT mocks!
 from agent_framework._workflows._events import (
    ExecutorCompletedEvent,
    ExecutorFailedEvent,
    ExecutorInvokedEvent,
    WorkflowErrorDetails,
 )
+from agent_framework.orchestrations import ConcurrentBuilder, SequentialBuilder

 from agent_framework_devui._discovery import EntityDiscovery
 from agent_framework_devui._executor import AgentFrameworkExecutor
 from agent_framework_devui._mapper import MessageMapper
 from agent_framework_devui.models._openai_custom import AgentFrameworkRequest

+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+
 # =============================================================================
 # Mock Chat Clients (from core tests pattern)
 # =============================================================================
@@ -92,7 +90,6 @@ class MockChatClient:
            yield ChatResponseUpdate(contents=[Content.from_text(text="test streaming response")], role="assistant")


-@use_chat_middleware
 class MockBaseChatClient(BaseChatClient[TOptions_co], Generic[TOptions_co]):
    """Full BaseChatClient mock with middleware support.

@@ -109,27 +106,27 @@ class MockBaseChatClient(BaseChatClient[TOptions_co], Generic[TOptions_co]):
        self.received_messages: list[list[ChatMessage]] = []

    @override
-    async def _inner_get_response(
+    def _inner_get_response(
        self,
        *,
-        messages: MutableSequence[ChatMessage],
-        options: dict[str, Any],
+        messages: Sequence[ChatMessage],
+        stream: bool,
+        options: Mapping[str, Any],
        **kwargs: Any,
-    ) -> ChatResponse:
-        self.call_count += 1
-        self.received_messages.append(list(messages))
-        if self.run_responses:
-            return self.run_responses.pop(0)
-        return ChatResponse(messages=ChatMessage("assistant", ["Mock response from ChatAgent"]))
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        if stream:
+            return self._build_response_stream(self._stream_impl(messages))

-    @override
-    async def _inner_get_streaming_response(
-        self,
-        *,
-        messages: MutableSequence[ChatMessage],
-        options: dict[str, Any],
-        **kwargs: Any,
-    ) -> AsyncIterable[ChatResponseUpdate]:
+        async def _get() -> ChatResponse:
+            self.call_count += 1
+            self.received_messages.append(list(messages))
+            if self.run_responses:
+                return self.run_responses.pop(0)
+            return ChatResponse(messages=ChatMessage("assistant", ["Mock response from ChatAgent"]))
+
+        return _get()
+
+    async def _stream_impl(self, messages: Sequence[ChatMessage]) -> AsyncIterable[ChatResponseUpdate]:
        self.call_count += 1
        self.received_messages.append(list(messages))
        if self.streaming_responses:
@@ -162,7 +159,20 @@ class MockAgent(BaseAgent):
        self.streaming_chunks = streaming_chunks or [response_text]
        self.call_count = 0

-    async def run(
+    def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        stream: bool = False,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+        self.call_count += 1
+        if stream:
+            return self._run_stream(messages=messages, thread=thread, **kwargs)
+        return self._run(messages=messages, thread=thread, **kwargs)
+
+    async def _run(
        self,
        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
        *,
@@ -172,16 +182,20 @@ class MockAgent(BaseAgent):
        self.call_count += 1
        return AgentResponse(messages=[ChatMessage("assistant", [Content.from_text(text=self.response_text)])])

-    async def run_stream(
+    def _run_stream(
        self,
        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
        *,
        thread: AgentThread | None = None,
        **kwargs: Any,
-    ) -> AsyncIterable[AgentResponseUpdate]:
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
        self.call_count += 1
-        for chunk in self.streaming_chunks:
-            yield AgentResponseUpdate(contents=[Content.from_text(text=chunk)], role="assistant")
+
+        async def _iter():
+            for chunk in self.streaming_chunks:
+                yield AgentResponseUpdate(contents=[Content.from_text(text=chunk)], role="assistant")
+
+        return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)


 class MockToolCallingAgent(BaseAgent):
@@ -191,115 +205,87 @@ class MockToolCallingAgent(BaseAgent):
        super().__init__(**kwargs)
        self.call_count = 0

-    async def run(
+    def run(
+        self,
+        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
+        *,
+        stream: bool = False,
+        thread: AgentThread | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+        self.call_count += 1
+        if stream:
+            return self._run_stream(messages=messages, thread=thread, **kwargs)
+        return self._run(messages=messages, thread=thread, **kwargs)
+
+    async def _run(
        self,
        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
        *,
        thread: AgentThread | None = None,
        **kwargs: Any,
    ) -> AgentResponse:
-        self.call_count += 1
        return AgentResponse(messages=[ChatMessage("assistant", ["done"])])

-    async def run_stream(
+    def _run_stream(
        self,
        messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
        *,
        thread: AgentThread | None = None,
        **kwargs: Any,
-    ) -> AsyncIterable[AgentResponseUpdate]:
-        self.call_count += 1
-        # First: text
-        yield AgentResponseUpdate(
-            contents=[Content.from_text(text="Let me search for that...")],
-            role="assistant",
-        )
-        # Second: tool call
-        yield AgentResponseUpdate(
-            contents=[
-                Content.from_function_call(
-                    call_id="call_123",
-                    name="search",
-                    arguments={"query": "weather"},
-                )
-            ],
-            role="assistant",
-        )
-        # Third: tool result
-        yield AgentResponseUpdate(
-            contents=[
-                Content.from_function_result(
-                    call_id="call_123",
-                    result={"temperature": 72, "condition": "sunny"},
-                )
-            ],
-            role="tool",
-        )
-        # Fourth: final text
-        yield AgentResponseUpdate(
-            contents=[Content.from_text(text="The weather is sunny, 72°F.")],
-            role="assistant",
-        )
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
+        async def _iter() -> AsyncIterable[AgentResponseUpdate]:
+            # First: text
+            yield AgentResponseUpdate(
+                contents=[Content.from_text(text="Let me search for that...")],
+                role="assistant",
+            )
+            # Second: tool call
+            yield AgentResponseUpdate(
+                contents=[
+                    Content.from_function_call(
+                        call_id="call_123",
+                        name="search",
+                        arguments={"query": "weather"},
+                    )
+                ],
+                role="assistant",
+            )
+            # Third: tool result
+            yield AgentResponseUpdate(
+                contents=[
+                    Content.from_function_result(
+                        call_id="call_123",
+                        result={"temperature": 72, "condition": "sunny"},
+                    )
+                ],
+                role="tool",
+            )
+            # Fourth: final text
+            yield AgentResponseUpdate(
+                contents=[Content.from_text(text="The weather is sunny, 72°F.")],
+                role="assistant",
+            )
+
+        return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)


 # =============================================================================
-# Factory Functions for Test Data
+# Helper Functions for Test Data Creation
 # =============================================================================


-def create_mapper() -> MessageMapper:
-    """Create a fresh MessageMapper."""
-    return MessageMapper()
-
-
-def create_test_request(
-    entity_id: str = "test_agent",
-    input_text: str = "Test input",
-    stream: bool = True,
-) -> AgentFrameworkRequest:
-    """Create a standard test request."""
-    return AgentFrameworkRequest(
-        metadata={"entity_id": entity_id},
-        input=input_text,
-        stream=stream,
-    )
-
-
-def create_mock_chat_client() -> MockChatClient:
-    """Create a mock chat client."""
-    return MockChatClient()
-
-
-def create_mock_base_chat_client() -> MockBaseChatClient:
-    """Create a mock BaseChatClient."""
-    return MockBaseChatClient()
-
-
-def create_mock_agent(
-    id: str = "test_agent",
-    name: str = "TestAgent",
-    response_text: str = "Mock agent response",
-) -> MockAgent:
-    """Create a mock agent."""
-    return MockAgent(id=id, name=name, response_text=response_text)
-
-
-def create_mock_tool_agent(id: str = "tool_agent", name: str = "ToolAgent") -> MockToolCallingAgent:
-    """Create a mock agent that simulates tool calls."""
-    return MockToolCallingAgent(id=id, name=name)
-
-
-def create_agent_run_response(text: str = "Test response") -> AgentResponse:
+def _create_agent_run_response(text: str = "Test response") -> AgentResponse:
    """Create an AgentResponse with the given text."""
    return AgentResponse(messages=[ChatMessage("assistant", [Content.from_text(text=text)])])


-def create_agent_executor_response(
+def _create_agent_executor_response(
    executor_id: str = "test_executor",
    response_text: str = "Executor response",
 ) -> AgentExecutorResponse:
    """Create an AgentExecutorResponse - the type that's nested in ExecutorCompletedEvent.data."""
-    agent_response = create_agent_run_response(response_text)
+    agent_response = _create_agent_run_response(response_text)
    return AgentExecutorResponse(
        executor_id=executor_id,
        agent_response=agent_response,
@@ -310,6 +296,21 @@ def create_agent_executor_response(
    )


+# =============================================================================
+# Public Factory Functions (for direct import in tests)
+# =============================================================================
+
+
+def create_agent_run_response(text: str = "Test response") -> AgentResponse:
+    """Create an AgentResponse with the given text."""
+    return _create_agent_run_response(text)
+
+
+def create_executor_invoked_event(executor_id: str = "test_executor") -> ExecutorInvokedEvent:
+    """Create an ExecutorInvokedEvent."""
+    return ExecutorInvokedEvent(executor_id=executor_id)
+
+
 def create_executor_completed_event(
    executor_id: str = "test_executor",
    with_agent_response: bool = True,
@@ -320,15 +321,10 @@ def create_executor_completed_event(
    ExecutorCompletedEvent.data contains AgentExecutorResponse which contains
    AgentResponse and ChatMessage objects (SerializationMixin, not Pydantic).
    """
-    data = create_agent_executor_response(executor_id) if with_agent_response else {"simple": "dict"}
+    data = _create_agent_executor_response(executor_id) if with_agent_response else {"simple": "dict"}
    return ExecutorCompletedEvent(executor_id=executor_id, data=data)


-def create_executor_invoked_event(executor_id: str = "test_executor") -> ExecutorInvokedEvent:
-    """Create an ExecutorInvokedEvent."""
-    return ExecutorInvokedEvent(executor_id=executor_id)
-
-
 def create_executor_failed_event(
    executor_id: str = "test_executor",
    error_message: str = "Test error",
@@ -339,11 +335,97 @@ def create_executor_failed_event(


 # =============================================================================
-# Workflow Setup Helpers (async factory functions)
+# Pytest Fixtures
 # =============================================================================


-async def create_executor_with_real_agent() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient]:
+@pytest.fixture
+def mapper() -> MessageMapper:
+    """Create a fresh MessageMapper for each test."""
+    return MessageMapper()
+
+
+@pytest.fixture
+def test_request() -> AgentFrameworkRequest:
+    """Create a standard test request."""
+    return AgentFrameworkRequest(
+        metadata={"entity_id": "test_agent"},
+        input="Test input",
+        stream=True,
+    )
+
+
+@pytest.fixture
+def mock_chat_client() -> MockChatClient:
+    """Create a mock chat client."""
+    return MockChatClient()
+
+
+@pytest.fixture
+def mock_base_chat_client() -> MockBaseChatClient:
+    """Create a mock BaseChatClient."""
+    return MockBaseChatClient()
+
+
+@pytest.fixture
+def mock_agent() -> MockAgent:
+    """Create a mock agent."""
+    return MockAgent(id="test_agent", name="TestAgent", response_text="Mock agent response")
+
+
+@pytest.fixture
+def mock_tool_agent() -> MockToolCallingAgent:
+    """Create a mock agent that simulates tool calls."""
+    return MockToolCallingAgent(id="tool_agent", name="ToolAgent")
+
+
+@pytest.fixture
+def agent_run_response() -> AgentResponse:
+    """Create an AgentResponse with default text."""
+    return _create_agent_run_response()
+
+
+@pytest.fixture
+def executor_completed_event() -> ExecutorCompletedEvent:
+    """Create an ExecutorCompletedEvent with realistic nested data.
+
+    This creates the exact data structure that caused the serialization bug:
+    ExecutorCompletedEvent.data contains AgentExecutorResponse which contains
+    AgentResponse and ChatMessage objects (SerializationMixin, not Pydantic).
+    """
+    data = _create_agent_executor_response("test_executor")
+    return ExecutorCompletedEvent(executor_id="test_executor", data=data)
+
+
+@pytest.fixture
+def executor_invoked_event() -> ExecutorInvokedEvent:
+    """Create an ExecutorInvokedEvent."""
+    return ExecutorInvokedEvent(executor_id="test_executor")
+
+
+@pytest.fixture
+def executor_failed_event() -> ExecutorFailedEvent:
+    """Create an ExecutorFailedEvent."""
+    details = WorkflowErrorDetails(error_type="TestError", message="Test error")
+    return ExecutorFailedEvent(executor_id="test_executor", details=details)
+
+
+@pytest.fixture
+def test_entities_dir() -> str:
+    """Use the samples directory which has proper entity structure."""
+    current_dir = Path(__file__).parent
+    # Navigate to python/samples/getting_started/devui
+    samples_dir = current_dir.parent.parent.parent.parent / "samples" / "getting_started" / "devui"
+    return str(samples_dir.resolve())
+
+
+# =============================================================================
+# Async Fixtures for Executor/Workflow Setup
+# =============================================================================
+
+
+@pytest_asyncio.fixture
+async def executor_with_real_agent() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient]:
    """Create an executor with a REAL ChatAgent using mock chat client.

    This tests the full execution pipeline:
@@ -375,7 +457,8 @@ async def create_executor_with_real_agent() -> tuple[AgentFrameworkExecutor, str
    return executor, entity_info.id, mock_client


-async def create_sequential_workflow() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient, Any]:
+@pytest_asyncio.fixture
+async def sequential_workflow() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient, Any]:
    """Create a realistic sequential workflow (Writer -> Reviewer).

    This provides a reusable multi-agent workflow that:
@@ -418,7 +501,8 @@ async def create_sequential_workflow() -> tuple[AgentFrameworkExecutor, str, Moc
    return executor, entity_info.id, mock_client, workflow


-async def create_concurrent_workflow() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient, Any]:
+@pytest_asyncio.fixture
+async def concurrent_workflow() -> tuple[AgentFrameworkExecutor, str, MockBaseChatClient, Any]:
    """Create a realistic concurrent workflow (Researcher | Analyst | Summarizer).

    This provides a reusable fan-out/fan-in workflow that:
@@ -338,7 +338,7 @@ class TestIntegration:
        checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)

        # Set build-time storage (equivalent to .with_checkpointing() at build time)
-        # Note: In production, DevUI uses runtime injection via run_stream() parameter
+        # Note: In production, DevUI uses runtime injection via run(stream=True) parameter
        if hasattr(test_workflow, "_runner") and hasattr(test_workflow._runner, "context"):
            test_workflow._runner.context._checkpoint_storage = checkpoint_storage

@@ -406,7 +406,7 @@ class TestIntegration:
        3. Framework automatically saves checkpoint to our storage
        4. Checkpoint is accessible via manager for UI to list/resume

-        Note: In production, DevUI passes checkpoint_storage to run_stream() as runtime parameter.
+        Note: In production, DevUI passes checkpoint_storage to run(stream=True) as runtime parameter.
        This test uses build-time injection to verify framework's checkpoint auto-save behavior.
        """
        entity_id = "test_entity"
@@ -427,7 +427,7 @@ class TestIntegration:

        # Run workflow until it reaches IDLE_WITH_PENDING_REQUESTS (after checkpoint is created)
        saw_request_event = False
-        async for event in test_workflow.run_stream(WorkflowTestData(value="test")):
+        async for event in test_workflow.run(WorkflowTestData(value="test"), stream=True):
            if isinstance(event, RequestInfoEvent):
                saw_request_event = True
            # Wait for IDLE_WITH_PENDING_REQUESTS status (comes after checkpoint creation)
@@ -33,10 +33,18 @@ class MockAgent:
        self.cleanup_called = False
        self.async_cleanup_called = False

-    async def run_stream(self, messages=None, *, thread=None, **kwargs):
-        """Mock streaming run method."""
-        yield AgentResponse(
-            messages=[ChatMessage("assistant", [Content.from_text(text="Test response")])],
+    async def run(self, messages=None, *, stream: bool = False, thread=None, **kwargs):
+        """Mock run method with streaming support."""
+        if stream:
+
+            async def _stream():
+                yield AgentResponse(
+                    messages=[ChatMessage(role="assistant", contents=[Content.from_text(text="Test response")])],
+                )
+
+            return _stream()
+        return AgentResponse(
+            messages=[ChatMessage(role="assistant", contents=[Content.from_text(text="Test response")])],
        )


@@ -277,9 +285,16 @@ class TestAgent:
    name = "Test Agent"
    description = "Test agent with cleanup"

-    async def run_stream(self, messages=None, *, thread=None, **kwargs):
-        yield AgentResponse(
-            messages=[ChatMessage("assistant", [Content.from_text(text="Test")])],
+    async def run(self, messages=None, *, stream: bool = False, thread=None, **kwargs):
+        if stream:
+            async def _stream():
+                yield AgentResponse(
+                    messages=[ChatMessage(role="assistant", content=[Content.from_text(text="Test")])],
+                    inner_messages=[],
+                )
+            return _stream()
+        return AgentResponse(
+            messages=[ChatMessage(role="assistant", content=[Content.from_text(text="Test")])],
            inner_messages=[],
        )

@@ -216,7 +216,7 @@ async def test_list_items_converts_function_calls():

    # Simulate messages from agent execution with function calls
    messages = [
-        ChatMessage("user", [{"type": "text", "text": "What's the weather in SF?"}]),
+        ChatMessage(role="user", contents=[{"type": "text", "text": "What's the weather in SF?"}]),
        ChatMessage(
            role="assistant",
            contents=[
@@ -238,7 +238,7 @@ async def test_list_items_converts_function_calls():
                }
            ],
        ),
-        ChatMessage("assistant", [{"type": "text", "text": "The weather is sunny, 65°F"}]),
+        ChatMessage(role="assistant", contents=[{"type": "text", "text": "The weather is sunny, 65°F"}]),
    ]

    # Add messages to thread
@@ -6,19 +6,9 @@ import asyncio
 import tempfile
 from pathlib import Path

-import pytest
-
 from agent_framework_devui._discovery import EntityDiscovery

-
-@pytest.fixture
-def test_entities_dir():
-    """Use the samples directory which has proper entity structure."""
-    # Get the samples directory from the main python samples folder
-    current_dir = Path(__file__).parent
-    # Navigate to python/samples/getting_started/devui
-    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
-    return str(samples_dir.resolve())
+# Note: test_entities_dir fixture is provided by conftest.py


 async def test_discover_agents(test_entities_dir):
@@ -89,7 +79,7 @@ from agent_framework import AgentResponse, AgentThread, ChatMessage, Role, Conte
 class NonStreamingAgent:
    id = "non_streaming"
    name = "Non-Streaming Agent"
-    description = "Agent without run_stream"
+    description = "Agent with run() method"

    async def run(self, messages=None, *, thread=None, **kwargs):
        return AgentResponse(
@@ -125,7 +115,6 @@ agent = NonStreamingAgent()
        enriched = discovery.get_entity_info(entity.id)
        assert enriched.type == "agent"  # Now correctly identified
        assert enriched.name == "Non-Streaming Agent"
-        assert not enriched.metadata.get("has_run_stream")


 async def test_lazy_loading():
@@ -210,7 +199,7 @@ class TestAgent:

    async def run(self, messages=None, *, thread=None, **kwargs):
        return AgentResponse(
-            messages=[ChatMessage("assistant", [Content.from_text(text="test")])],
+            messages=[ChatMessage(role="assistant", contents=[Content.from_text(text="test")])],
            response_id="test"
        )

@@ -342,7 +331,7 @@ class WeatherAgent:
    name = "Weather Agent"
    description = "Gets weather information"

-    def run_stream(self, input_str):
+    def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
        return f"Weather in {input_str}"
 """)

@@ -15,16 +15,10 @@ from pathlib import Path
 from typing import Any

 import pytest
-import pytest_asyncio
 from agent_framework import AgentExecutor, ChatAgent, FunctionExecutor, WorkflowBuilder

-# Import test utilities
-from test_helpers import (
-    MockBaseChatClient,
-    create_concurrent_workflow,
-    create_executor_with_real_agent,
-    create_sequential_workflow,
-)
+# Import mock classes from conftest for direct use in some tests
+from conftest import MockBaseChatClient

 from agent_framework_devui._discovery import EntityDiscovery
 from agent_framework_devui._executor import AgentFrameworkExecutor, EntityNotFoundError
@@ -32,38 +26,10 @@ from agent_framework_devui._mapper import MessageMapper
 from agent_framework_devui.models._openai_custom import AgentFrameworkRequest

 # =============================================================================
-# Local Fixtures (async factory-based)
+# Local Fixtures (module-specific)
 # =============================================================================


-@pytest_asyncio.fixture
-async def executor_with_real_agent():
-    """Create an executor with a REAL ChatAgent using mock chat client."""
-    return await create_executor_with_real_agent()
-
-
-@pytest_asyncio.fixture
-async def sequential_workflow_fixture():
-    """Create a realistic sequential workflow (Writer -> Reviewer)."""
-    return await create_sequential_workflow()
-
-
-@pytest_asyncio.fixture
-async def concurrent_workflow_fixture():
-    """Create a realistic concurrent workflow (Researcher | Analyst | Summarizer)."""
-    return await create_concurrent_workflow()
-
-
-@pytest.fixture
-def test_entities_dir():
-    """Use the samples directory which has proper entity structure."""
-    # Get the samples directory from the main python samples folder
-    current_dir = Path(__file__).parent
-    # Navigate to python/samples/getting_started/devui
-    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
-    return str(samples_dir.resolve())
-
-
@pytest.fixture
 async def executor(test_entities_dir):
    """Create configured executor."""
@@ -419,9 +385,9 @@ async def test_request_extracts_entity_id_from_metadata(executor):


@pytest.mark.asyncio
-async def test_executor_get_start_executor_message_types(sequential_workflow_fixture):
+async def test_executor_get_start_executor_message_types(sequential_workflow):
    """Test _get_start_executor_message_types with real workflow."""
-    executor, _entity_id, _mock_client, workflow = sequential_workflow_fixture
+    executor, _entity_id, _mock_client, workflow = sequential_workflow

    start_exec, message_types = executor._get_start_executor_message_types(workflow)

@@ -493,11 +459,11 @@ async def test_executor_parse_raw_string_for_string_workflow():


@pytest.mark.asyncio
-async def test_executor_parse_converts_to_chat_message_for_sequential_workflow(sequential_workflow_fixture):
+async def test_executor_parse_converts_to_chat_message_for_sequential_workflow(sequential_workflow):
    """Sequential workflows convert string input to ChatMessage."""
    from agent_framework import ChatMessage

-    executor, _entity_id, _mock_client, workflow = sequential_workflow_fixture
+    executor, _entity_id, _mock_client, workflow = sequential_workflow

    # Sequential workflows expect ChatMessage, so raw string becomes ChatMessage
    parsed = executor._parse_raw_workflow_input(workflow, "hello")
@@ -564,23 +530,36 @@ def test_extract_workflow_hil_responses_handles_stringified_json():
    assert executor._extract_workflow_hil_responses({"email": "test"}) is None


-async def test_executor_handles_non_streaming_agent():
-    """Test executor can handle agents with only run() method (no run_stream)."""
-    from agent_framework import AgentResponse, AgentThread, ChatMessage, Content
+async def test_executor_handles_streaming_agent():
+    """Test executor handles agents with run(stream=True) method."""
+    from agent_framework import AgentResponse, AgentResponseUpdate, AgentThread, ChatMessage, Content

-    class NonStreamingAgent:
-        """Agent with only run() method - does NOT satisfy full AgentProtocol."""
+    class StreamingAgent:
+        """Agent with run() method supporting stream parameter."""

-        id = "non_streaming_test"
-        name = "Non-Streaming Test Agent"
-        description = "Test agent without run_stream()"
+        id = "streaming_test"
+        name = "Streaming Test Agent"
+        description = "Test agent with run(stream=True)"

-        async def run(self, messages=None, *, thread=None, **kwargs):
+        def run(self, messages=None, *, stream=False, thread=None, **kwargs):
+            if stream:
+                # Return an async generator for streaming
+                return self._stream_impl(messages)
+            # Return awaitable for non-streaming
+            return self._run_impl(messages)
+
+        async def _run_impl(self, messages):
            return AgentResponse(
-                messages=[ChatMessage("assistant", [Content.from_text(text=f"Processed: {messages}")])],
+                messages=[ChatMessage(role="assistant", contents=[Content.from_text(text=f"Processed: {messages}")])],
                response_id="test_123",
            )

+        async def _stream_impl(self, messages):
+            yield AgentResponseUpdate(
+                contents=[Content.from_text(text=f"Processed: {messages}")],
+                role="assistant",
+            )
+
        def get_new_thread(self, **kwargs):
            return AgentThread()

@@ -589,11 +568,11 @@ async def test_executor_handles_non_streaming_agent():
    mapper = MessageMapper()
    executor = AgentFrameworkExecutor(discovery, mapper)

-    agent = NonStreamingAgent()
+    agent = StreamingAgent()
    entity_info = await discovery.create_entity_info_from_object(agent, source="test")
    discovery.register_entity(entity_info.id, entity_info, agent)

-    # Execute non-streaming agent (use metadata.entity_id for routing)
+    # Execute streaming agent (use metadata.entity_id for routing)
    request = AgentFrameworkRequest(
        metadata={"entity_id": entity_info.id},
        input="hello",
@@ -604,7 +583,7 @@ async def test_executor_handles_non_streaming_agent():
    async for event in executor.execute_streaming(request):
        events.append(event)

-    # Should get events even though agent doesn't stream
+    # Should get events from streaming agent
    assert len(events) > 0
    text_events = [e for e in events if hasattr(e, "type") and e.type == "response.output_text.delta"]
    assert len(text_events) > 0
@@ -617,13 +596,13 @@ async def test_executor_handles_non_streaming_agent():


@pytest.mark.asyncio
-async def test_full_pipeline_sequential_workflow(sequential_workflow_fixture):
+async def test_full_pipeline_sequential_workflow(sequential_workflow):
    """Test SequentialBuilder workflow full pipeline with JSON serialization.

-    Uses the shared sequential_workflow_fixture (Writer → Reviewer) from conftest.
+    Uses the shared sequential_workflow fixture (Writer → Reviewer) from conftest.
    Tests that all events can be JSON serialized for SSE streaming.
    """
-    executor, entity_id, mock_client, _workflow = sequential_workflow_fixture
+    executor, entity_id, mock_client, _workflow = sequential_workflow

    request = AgentFrameworkRequest(
        metadata={"entity_id": entity_id},
@@ -652,13 +631,13 @@ async def test_full_pipeline_sequential_workflow(sequential_workflow_fixture):


@pytest.mark.asyncio
-async def test_full_pipeline_concurrent_workflow(concurrent_workflow_fixture):
+async def test_full_pipeline_concurrent_workflow(concurrent_workflow):
    """Test ConcurrentBuilder workflow full pipeline with JSON serialization.

-    Uses the shared concurrent_workflow_fixture (Researcher | Analyst | Summarizer) from conftest.
+    Uses the shared concurrent_workflow fixture (Researcher | Analyst | Summarizer) from conftest.
    Tests fan-out/fan-in pattern with parallel agent execution.
    """
-    executor, entity_id, mock_client, _workflow = concurrent_workflow_fixture
+    executor, entity_id, mock_client, _workflow = concurrent_workflow

    request = AgentFrameworkRequest(
        metadata={"entity_id": entity_id},
@@ -769,9 +748,13 @@ class StreamingAgent:
    name = "Streaming Test Agent"
    description = "Test agent for streaming"

-    async def run_stream(self, input_str):
-        for i, word in enumerate(f"Processing {input_str}".split()):
-            yield f"word_{i}: {word} "
+    async def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
+        if stream:
+            async def _stream():
+                for i, word in enumerate(f"Processing {input_str}".split()):
+                    yield f"word_{i}: {word} "
+            return _stream()
+        return f"Processing {input_str}"
 """)

            discovery = EntityDiscovery(str(temp_path))
@@ -24,14 +24,12 @@ from agent_framework._workflows._events import (
    WorkflowStatusEvent,
 )

-# Import test utilities
-from test_helpers import (
+# Import factory functions from conftest for parameterized test data creation
+from conftest import (
    create_agent_run_response,
    create_executor_completed_event,
    create_executor_failed_event,
    create_executor_invoked_event,
-    create_mapper,
-    create_test_request,
 )

 from agent_framework_devui._mapper import MessageMapper
@@ -42,21 +40,7 @@ from agent_framework_devui.models._openai_custom import (
    AgentStartedEvent,
 )

-# =============================================================================
-# Local Fixtures (to replace conftest.py fixtures)
-# =============================================================================
-
-
-@pytest.fixture
-def mapper() -> MessageMapper:
-    """Create a fresh MessageMapper for each test."""
-    return create_mapper()
-
-
-@pytest.fixture
-def test_request() -> AgentFrameworkRequest:
-    """Create a standard test request."""
-    return create_test_request()
+# Note: mapper and test_request fixtures are provided by conftest.py


 # =============================================================================
@@ -602,8 +586,8 @@ async def test_workflow_output_event_with_list_data(mapper: MessageMapper, test_

    # Sequential/Concurrent workflows often output list[ChatMessage]
    messages = [
-        ChatMessage("user", [Content.from_text(text="Hello")]),
-        ChatMessage("assistant", [Content.from_text(text="World")]),
+        ChatMessage(role="user", contents=[Content.from_text(text="Hello")]),
+        ChatMessage(role="assistant", contents=[Content.from_text(text="World")]),
    ]
    event = WorkflowOutputEvent(data=messages, executor_id="complete")
    events = await mapper.convert_event(event, test_request)
@@ -86,9 +86,8 @@ class TestMultimodalWorkflowInput:
        assert result.contents[1].media_type == "image/png"
        assert result.contents[1].uri == TEST_IMAGE_DATA_URI

-    def test_parse_workflow_input_handles_json_string_with_multimodal(self):
+    async def test_parse_workflow_input_handles_json_string_with_multimodal(self):
        """Test that _parse_workflow_input correctly handles JSON string with multimodal content."""
-        import asyncio

        from agent_framework import ChatMessage

@@ -113,7 +112,7 @@ class TestMultimodalWorkflowInput:
        mock_workflow = MagicMock()

        # Parse the input
-        result = asyncio.run(executor._parse_workflow_input(mock_workflow, json_string_input))
+        result = await executor._parse_workflow_input(mock_workflow, json_string_input)

        # Verify result is ChatMessage with multimodal content
        assert isinstance(result, ChatMessage), f"Expected ChatMessage, got {type(result)}"
@@ -127,9 +126,8 @@ class TestMultimodalWorkflowInput:
        assert result.contents[1].type == "data"
        assert result.contents[1].media_type == "image/png"

-    def test_parse_workflow_input_still_handles_simple_dict(self):
+    async def test_parse_workflow_input_still_handles_simple_dict(self):
        """Test that simple dict input still works (backward compatibility)."""
-        import asyncio

        from agent_framework import ChatMessage

@@ -148,7 +146,7 @@ class TestMultimodalWorkflowInput:
        mock_workflow.get_start_executor.return_value = mock_executor

        # Parse the input
-        result = asyncio.run(executor._parse_workflow_input(mock_workflow, json_string_input))
+        result = await executor._parse_workflow_input(mock_workflow, json_string_input)

        # Result should be ChatMessage (from _parse_structured_workflow_input)
        assert isinstance(result, ChatMessage), f"Expected ChatMessage, got {type(result)}"
@@ -23,14 +23,7 @@ class _StubExecutor:
            self._handlers = dict(handlers)


-@pytest.fixture
-def test_entities_dir():
-    """Use the samples directory which has proper entity structure."""
-    # Get the samples directory from the main python samples folder
-    current_dir = Path(__file__).parent
-    # Navigate to python/samples/getting_started/devui
-    samples_dir = current_dir.parent.parent.parent / "samples" / "getting_started" / "devui"
-    return str(samples_dir.resolve())
+# Note: test_entities_dir fixture is provided by conftest.py


 async def test_server_health_endpoint(test_entities_dir):
@@ -159,6 +152,7 @@ async def test_credential_cleanup() -> None:
    mock_client = Mock()
    mock_client.async_credential = mock_credential
    mock_client.model_id = "test-model"
+    mock_client.function_invocation_configuration = None

    # Create agent with mock client
    agent = ChatAgent(name="TestAgent", chat_client=mock_client, instructions="Test agent")
@@ -191,6 +185,7 @@ async def test_credential_cleanup_error_handling() -> None:
    mock_client = Mock()
    mock_client.async_credential = mock_credential
    mock_client.model_id = "test-model"
+    mock_client.function_invocation_configuration = None

    # Create agent with mock client
    agent = ChatAgent(name="TestAgent", chat_client=mock_client, instructions="Test agent")
@@ -225,6 +220,7 @@ async def test_multiple_credential_attributes() -> None:
    mock_client.credential = mock_cred1
    mock_client.async_credential = mock_cred2
    mock_client.model_id = "test-model"
+    mock_client.function_invocation_configuration = None

    # Create agent with mock client
    agent = ChatAgent(name="TestAgent", chat_client=mock_client, instructions="Test agent")
@@ -346,7 +342,7 @@ class WeatherAgent:
    name = "Weather Agent"
    description = "Gets weather information"

-    def run_stream(self, input_str):
+    def run(self, input_str, *, stream: bool = False, thread=None, **kwargs):
        return f"Weather in {input_str} is sunny"
 """)