Merge remote-tracking branch 'origin/main' into copilot/fix-azure-functions-worker-crashes

# Conflicts:
#	.github/workflows/python-integration-tests.yml
#	.github/workflows/python-merge-tests.yml

python/packages/azurefunctions/agent_framework_azurefunctions/_app.py

@@ -14,6 +14,7 @@ import logging
 import re
 import uuid
 from collections.abc import Callable, Mapping
+from copy import deepcopy
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, TypeVar, cast
@@ -21,6 +22,7 @@ from typing import TYPE_CHECKING, Any, TypeVar, cast
 import azure.durable_functions as df
 import azure.functions as func
 from agent_framework import AgentExecutor, SupportsAgentRun, Workflow, WorkflowEvent
+from agent_framework._workflows._runner_context import YieldOutputEventType
 from agent_framework_durabletask import (
     DEFAULT_MAX_POLL_RETRIES,
     DEFAULT_POLL_INTERVAL_SECONDS,
@@ -44,7 +46,7 @@ from ._context import CapturingRunnerContext
 from ._entities import create_agent_entity
 from ._errors import IncomingRequestError
 from ._orchestration import AgentOrchestrationContextType, AgentTask, AzureFunctionsAgentExecutor
-from ._serialization import deserialize_value, serialize_value
+from ._serialization import deserialize_value, serialize_value, strip_pickle_markers
 from ._workflow import (
     SOURCE_HITL_RESPONSE,
     SOURCE_ORCHESTRATOR,
@@ -58,6 +60,11 @@ EntityHandler = Callable[[df.DurableEntityContext], None]
 HandlerT = TypeVar("HandlerT", bound=Callable[..., Any])
 
 
+def _create_state_snapshot(state: dict[str, Any]) -> dict[str, Any]:
+    """Create a deep copy of the deserialized state for later diffing."""
+    return deepcopy(state)
+
+
 @dataclass
 class AgentMetadata:
     """Metadata for a registered agent.
@@ -118,16 +125,17 @@ class AgentFunctionApp(DFAppBase):
 
     .. code-block:: python
 
-        from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient
+        from agent_framework.azure import AgentFunctionApp
+        from agent_framework.openai import OpenAIChatCompletionClient
 
         # Create agents with unique names
-        weather_agent = AzureOpenAIChatClient(...).as_agent(
+        weather_agent = OpenAIChatCompletionClient(...).as_agent(
             name="WeatherAgent",
             instructions="You are a helpful weather agent.",
             tools=[get_weather],
         )
 
-        math_agent = AzureOpenAIChatClient(...).as_agent(
+        math_agent = OpenAIChatCompletionClient(...).as_agent(
             name="MathAgent",
             instructions="You are a helpful math assistant.",
             tools=[calculate],
@@ -274,10 +282,14 @@ class AgentFunctionApp(DFAppBase):
             """
             from agent_framework._workflows._state import State
 
-            data = json.loads(inputData)
-            message_data = data["message"]
+            data_obj = json.loads(inputData)
+            if not isinstance(data_obj, dict):
+                raise ValueError("Activity inputData must decode to a JSON object")
+            data = cast(dict[str, Any], data_obj)
+
+            message_data = data.get("message")
             shared_state_snapshot = data.get("shared_state_snapshot", {})
-            source_executor_ids = data.get("source_executor_ids", [SOURCE_ORCHESTRATOR])
+            source_executor_ids = cast(list[str], data.get("source_executor_ids", [SOURCE_ORCHESTRATOR]))
 
             if not self.workflow:
                 raise RuntimeError("Workflow not initialized in AgentFunctionApp")
@@ -296,18 +308,35 @@ class AgentFunctionApp(DFAppBase):
             async def run() -> dict[str, Any]:
                 # Create runner context and shared state
                 runner_context = CapturingRunnerContext()
+                workflow = self.workflow
+
+                def classify_yielded_output(executor_id: str) -> YieldOutputEventType | None:
+                    if workflow is None:
+                        return "output"
+                    if workflow.is_terminal_executor(executor_id):
+                        return "output"
+                    if workflow.is_intermediate_executor(executor_id):
+                        return "intermediate"
+                    return None
+
+                runner_context.set_yield_output_classifier(classify_yielded_output)
                 shared_state = State()
 
                 # Deserialize shared state values to reconstruct dataclasses/Pydantic models
-                deserialized_state = {k: deserialize_value(v) for k, v in (shared_state_snapshot or {}).items()}
-                original_snapshot = dict(deserialized_state)
+                deserialized_state: dict[str, Any] = {
+                    str(k): deserialize_value(v) for k, v in shared_state_snapshot.items()
+                }
+                original_snapshot = _create_state_snapshot(deserialized_state)
                 shared_state.import_state(deserialized_state)
 
                 if is_hitl_response:
                     # Handle HITL response by calling the executor's @response_handler
+                    if not isinstance(message_data, dict):
+                        raise ValueError("HITL message payload must be a JSON object")
+
                     await execute_hitl_response_handler(
                         executor=executor,
-                        hitl_message=message_data,
+                        hitl_message=cast(dict[str, Any], message_data),
                         shared_state=shared_state,
                         runner_context=runner_context,
                     )
@@ -323,16 +352,17 @@ class AgentFunctionApp(DFAppBase):
                 # Commit pending state changes and export
                 shared_state.commit()
                 current_state = shared_state.export_state()
-                original_keys = set(original_snapshot.keys())
-                current_keys = set(current_state.keys())
+                original_keys: set[str] = set(original_snapshot.keys())
+                current_keys: set[str] = set(current_state.keys())
 
                 # Deleted = was in original, not in current
-                deletes = original_keys - current_keys
+                deletes: set[str] = original_keys - current_keys
 
                 # Updates = keys in current that are new or have different values
-                updates = {
-                    k: v for k, v in current_state.items() if k not in original_snapshot or original_snapshot[k] != v
-                }
+                updates: dict[str, Any] = {}
+                for key in current_keys:
+                    if key not in original_keys or current_state[key] != original_snapshot.get(key):
+                        updates[key] = current_state[key]
 
                 # Drain messages and events from runner context
                 sent_messages = await runner_context.drain_messages()
@@ -348,7 +378,7 @@ class AgentFunctionApp(DFAppBase):
                 pending_request_info_events = await runner_context.get_pending_request_info_events()
 
                 # Serialize pending request info events for orchestrator
-                serialized_pending_requests = []
+                serialized_pending_requests: list[dict[str, Any]] = []
                 for _request_id, event in pending_request_info_events.items():
                     serialized_pending_requests.append({
                         "request_id": event.request_id,
@@ -361,7 +391,7 @@ class AgentFunctionApp(DFAppBase):
                     })
 
                 # Serialize messages for JSON compatibility
-                serialized_sent_messages = []
+                serialized_sent_messages: list[dict[str, Any]] = []
                 for _source_id, msg_list in sent_messages.items():
                     for msg in msg_list:
                         serialized_sent_messages.append({
@@ -441,6 +471,9 @@ class AgentFunctionApp(DFAppBase):
         ) -> func.HttpResponse:
             """HTTP endpoint to get workflow status."""
             instance_id = req.route_params.get("instanceId")
+            if not instance_id:
+                return self._build_error_response("Instance ID is required", status_code=400)
+
             status = await client.get_status(instance_id)
 
             if not status:
@@ -457,17 +490,23 @@ class AgentFunctionApp(DFAppBase):
             }
 
             # Add pending HITL requests info if available
-            custom_status = status.custom_status or {}
-            if isinstance(custom_status, dict) and custom_status.get("pending_requests"):
+            if (
+                (custom_status := status.custom_status)
+                and isinstance(custom_status, dict)
+                and (pending_requests_dict := custom_status.get("pending_requests"))  # type: ignore
+                and isinstance(pending_requests_dict, dict)
+            ):
                 base_url = self._build_base_url(req.url)
-                pending_requests = []
-                for req_id, req_data in custom_status["pending_requests"].items():
+                pending_requests: list[dict[str, Any]] = []
+                for req_id, req_data in pending_requests_dict.items():  # type: ignore
+                    if not isinstance(req_data, dict):
+                        continue
                     pending_requests.append({
                         "requestId": req_id,
-                        "sourceExecutor": req_data.get("source_executor_id"),
-                        "requestData": req_data.get("data"),
-                        "requestType": req_data.get("request_type"),
-                        "responseType": req_data.get("response_type"),
+                        "sourceExecutor": req_data.get("source_executor_id"),  # type: ignore[reportUnknownMemberType]
+                        "requestData": req_data.get("data"),  # type: ignore[reportUnknownMemberType]
+                        "requestType": req_data.get("request_type"),  # type: ignore[reportUnknownMemberType]
+                        "responseType": req_data.get("response_type"),  # type: ignore[reportUnknownMemberType]
                         "respondUrl": f"{base_url}/api/workflow/respond/{instance_id}/{req_id}",
                     })
                 response["pendingHumanInputRequests"] = pending_requests
@@ -497,6 +536,10 @@ class AgentFunctionApp(DFAppBase):
             except ValueError:
                 return self._build_error_response("Request body must be valid JSON.")
 
+            # Sanitize untrusted HTTP input before it reaches pickle.loads().
+            # See strip_pickle_markers() docstring for details on the attack vector.
+            response_data = strip_pickle_markers(response_data)
+
             # Send the response as an external event
             # The request_id is used as the event name for correlation
             await client.raise_event(
@@ -515,6 +558,11 @@ class AgentFunctionApp(DFAppBase):
                 mimetype="application/json",
             )
 
+        # Ensure route handlers are registered (prevents unused function warnings)
+        _ = start_workflow_orchestration
+        _ = get_workflow_status
+        _ = send_hitl_response
+
     def _build_status_url(self, request_url: str, instance_id: str) -> str:
         """Build the status URL for a workflow instance."""
         base_url = self._build_base_url(request_url)

python/packages/azurefunctions/agent_framework_azurefunctions/_context.py

@@ -19,6 +19,7 @@ from agent_framework import (
     WorkflowEvent,
     WorkflowMessage,
 )
+from agent_framework._workflows._runner_context import YieldOutputClassifier, YieldOutputEventType
 from agent_framework._workflows._state import State
 
 
@@ -41,6 +42,7 @@ class CapturingRunnerContext(RunnerContext):
         self._pending_request_info_events: dict[str, WorkflowEvent[Any]] = {}
         self._workflow_id: str | None = None
         self._streaming: bool = False
+        self._yield_output_classifier: YieldOutputClassifier = lambda _executor_id: "output"
 
     # region Messaging
 
@@ -144,6 +146,14 @@ class CapturingRunnerContext(RunnerContext):
         """Check if streaming mode is enabled (always False in activity context)."""
         return self._streaming
 
+    def set_yield_output_classifier(self, classifier: YieldOutputClassifier) -> None:
+        """Set the classifier used by WorkflowContext.yield_output()."""
+        self._yield_output_classifier = classifier
+
+    def classify_yielded_output(self, executor_id: str) -> YieldOutputEventType | None:
+        """Classify an executor's yield_output payload as output, intermediate, or hidden."""
+        return self._yield_output_classifier(executor_id)
+
     # endregion Workflow Configuration
 
     # region Request Info Events

python/packages/azurefunctions/agent_framework_azurefunctions/_serialization.py

@@ -13,22 +13,29 @@ This module adds:
 - serialize_value / deserialize_value: convenience aliases for encode/decode
 - reconstruct_to_type: for HITL responses where external data (without type markers)
   needs to be reconstructed to a known type
-- _resolve_type: resolves 'module:class' type keys to Python types
+- resolve_type: resolves 'module:class' type keys to Python types
 """
 
 from __future__ import annotations
 
 import importlib
 import logging
+from contextlib import suppress
 from dataclasses import is_dataclass
-from typing import Any
+from typing import Any, cast
 
-from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value
+from agent_framework._workflows._checkpoint_encoding import (
+    _PICKLE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    _TYPE_MARKER,  # pyright: ignore[reportPrivateUsage]
+    decode_checkpoint_value,
+    encode_checkpoint_value,
+)
+from pydantic import BaseModel
 
 logger = logging.getLogger(__name__)
 
 
-def _resolve_type(type_key: str) -> type | None:
+def resolve_type(type_key: str) -> type | None:
     """Resolve a 'module:class' type key to its Python type.
 
     Args:
@@ -46,6 +53,41 @@ def _resolve_type(type_key: str) -> type | None:
         return None
 
 
+# ============================================================================
+# Pickle marker sanitization (security)
+# ============================================================================
+
+
+def strip_pickle_markers(data: Any) -> Any:
+    """Recursively strip pickle/type markers from untrusted data.
+
+    The core checkpoint encoding uses ``__pickled__`` and ``__type__`` markers to
+    roundtrip arbitrary Python objects via *pickle*.  If an attacker crafts an
+    HTTP payload that contains these markers, the data would flow into
+    ``pickle.loads()`` and enable **arbitrary code execution**.
+
+    This function walks the incoming data structure and replaces any ``dict``
+    that contains either marker key with ``None``, neutralising the attack
+    vector while leaving all other data untouched.
+
+    It **must** be called on every value that originates from an untrusted
+    source (e.g. ``req.get_json()``) *before* the value is passed to
+    ``deserialize_value`` / ``decode_checkpoint_value``.
+    """
+    if isinstance(data, dict):
+        if _PICKLE_MARKER in data or _TYPE_MARKER in data:
+            logger.debug("Stripped pickle/type markers from untrusted input.")
+            return None
+        typed_dict = cast(dict[str, Any], data)
+        return {k: strip_pickle_markers(v) for k, v in typed_dict.items()}
+
+    if isinstance(data, list):
+        typed_list = cast(list[Any], data)  # type: ignore[redundant-cast]
+        return [strip_pickle_markers(item) for item in typed_list]
+
+    return data
+
+
 # ============================================================================
 # Serialize / Deserialize
 # ============================================================================
@@ -108,32 +150,34 @@ def reconstruct_to_type(value: Any, target_type: type) -> Any:
     if value is None:
         return None
 
-    try:
+    with suppress(TypeError):
         if isinstance(value, target_type):
             return value
-    except TypeError:
-        pass
 
     if not isinstance(value, dict):
         return value
 
-    # Try decoding if data has pickle markers (from checkpoint encoding)
+    # Try decoding if data has pickle markers (from checkpoint encoding).
+    # NOTE: This function is general-purpose.  Callers that handle untrusted
+    # data (e.g. HITL responses) MUST call strip_pickle_markers() before
+    # passing data here.  See _deserialize_hitl_response in _workflow.py.
     decoded = deserialize_value(value)
     if not isinstance(decoded, dict):
         return decoded
 
     # Try Pydantic model validation (for unmarked dicts, e.g., external HITL data)
-    if hasattr(target_type, "model_validate"):
+    if issubclass(target_type, BaseModel):
         try:
             return target_type.model_validate(value)
         except Exception:
             logger.debug("Could not validate Pydantic model %s", target_type)
+            return value  # type: ignore[return-value]
 
     # Try dataclass construction (for unmarked dicts, e.g., external HITL data)
-    if is_dataclass(target_type) and isinstance(target_type, type):
+    if is_dataclass(target_type) and isinstance(target_type, type):  # type: ignore
         try:
             return target_type(**value)
         except Exception:
             logger.debug("Could not construct dataclass %s", target_type)
 
-    return value
+    return value  # type: ignore[return-value]

python/packages/azurefunctions/agent_framework_azurefunctions/_workflow.py

@@ -44,12 +44,13 @@ from agent_framework._workflows._edge import (
     SingleEdgeGroup,
     SwitchCaseEdgeGroup,
 )
+from agent_framework._workflows._state import State
 from agent_framework_durabletask import AgentSessionId, DurableAgentSession, DurableAIAgent
 from azure.durable_functions import DurableOrchestrationContext
 
 from ._context import CapturingRunnerContext
 from ._orchestration import AzureFunctionsAgentExecutor
-from ._serialization import _resolve_type, deserialize_value, reconstruct_to_type, serialize_value
+from ._serialization import deserialize_value, reconstruct_to_type, resolve_type, serialize_value, strip_pickle_markers
 
 logger = logging.getLogger(__name__)
 
@@ -148,7 +149,7 @@ def _evaluate_edge_condition_sync(edge: Edge, message: Any) -> bool:
         True if the edge should be traversed, False otherwise
     """
     # Access the internal condition directly since should_route is async
-    condition = edge._condition
+    condition = edge._condition  # pyright: ignore[reportPrivateUsage]
     if condition is None:
         return True
     result = condition(message)
@@ -239,11 +240,11 @@ def build_agent_executor_response(
     Returns:
         AgentExecutorResponse with reconstructed conversation
     """
-    final_text = response_text
+    final_text: str = response_text or ""
     if structured_response:
         final_text = json.dumps(structured_response)
 
-    assistant_message = Message(role="assistant", text=final_text)
+    assistant_message = Message(role="assistant", contents=[final_text])
 
     agent_response = AgentResponse(
         messages=[assistant_message],
@@ -254,7 +255,7 @@ def build_agent_executor_response(
     if isinstance(previous_message, AgentExecutorResponse) and previous_message.full_conversation:
         full_conversation.extend(previous_message.full_conversation)
     elif isinstance(previous_message, str):
-        full_conversation.append(Message(role="user", text=previous_message))
+        full_conversation.append(Message(role="user", contents=[previous_message]))
 
     full_conversation.append(assistant_message)
 
@@ -322,7 +323,8 @@ def _prepare_activity_task(
     activity_input_json = json.dumps(activity_input)
     # Use the prefixed activity name that matches the registered function
     activity_name = f"dafx-{executor_id}"
-    return context.call_activity(activity_name, activity_input_json)
+    orchestration_context: Any = context
+    return orchestration_context.call_activity(activity_name, activity_input_json)
 
 
 # ============================================================================
@@ -346,13 +348,16 @@ def _process_agent_response(
         ExecutorResult containing the processed response
     """
     response_text = agent_response.text if agent_response else None
-    structured_response = None
+    structured_response: dict[str, Any] | None = None
 
     if agent_response and agent_response.value is not None:
-        if hasattr(agent_response.value, "model_dump"):
-            structured_response = agent_response.value.model_dump()
+        model_dump = getattr(agent_response.value, "model_dump", None)
+        if callable(model_dump):
+            dumped = model_dump()
+            if isinstance(dumped, dict):
+                structured_response = dumped  # type: ignore[assignment]
         elif isinstance(agent_response.value, dict):
-            structured_response = agent_response.value
+            structured_response = agent_response.value  # type: ignore[assignment]
 
     output_message = build_agent_executor_response(
         executor_id=executor_id,
@@ -726,7 +731,7 @@ def run_workflow_orchestrator(
 
                 if winner == approval_task:
                     # Cancel the timeout
-                    timeout_task.cancel()
+                    timeout_task.cancel()  # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
 
                     # Get the response
                     raw_response = approval_task.result
@@ -756,7 +761,7 @@ def run_workflow_orchestrator(
                     )
                 else:
                     # Timeout occurred — cancel the dangling external event listener
-                    approval_task.cancel()
+                    approval_task.cancel()  # pyright: ignore[reportUnknownMemberType, reportAttributeAccessIssue]
                     logger.warning("HITL request %s timed out after %s hours", request_id, hitl_timeout_hours)
                     raise TimeoutError(
                         f"Human-in-the-loop request '{request_id}' timed out after {hitl_timeout_hours} hours."
@@ -864,7 +869,8 @@ def _extract_message_content(message: Any) -> str:
         # Extract text from the last message in the request
         message_content = message.messages[-1].text or ""
     elif isinstance(message, dict):
-        logger.warning("Unexpected dict message in _extract_message_content. Keys: %s", list(message.keys()))
+        key_names = list(message.keys())  # type: ignore[union-attr]
+        logger.warning("Unexpected dict message in _extract_message_content. Keys: %s", key_names)  # type: ignore
     elif isinstance(message, str):
         message_content = message
 
@@ -879,7 +885,7 @@ def _extract_message_content(message: Any) -> str:
 async def execute_hitl_response_handler(
     executor: Any,
     hitl_message: dict[str, Any],
-    shared_state: Any,
+    shared_state: State,
     runner_context: CapturingRunnerContext,
 ) -> None:
     """Execute a HITL response handler on an executor.
@@ -910,7 +916,7 @@ async def execute_hitl_response_handler(
     response = _deserialize_hitl_response(response_data, response_type_str)
 
     # Find the matching response handler
-    handler = executor._find_response_handler(original_request, response)
+    handler = executor._find_response_handler(original_request, response)  # pyright: ignore[reportPrivateUsage]
 
     if handler is None:
         logger.warning(
@@ -955,6 +961,13 @@ def _deserialize_hitl_response(response_data: Any, response_type_str: str | None
         type(response_data).__name__,
     )
 
+    if response_data is None:
+        return None
+
+    # Sanitize untrusted external input before deserialization.
+    # HITL response data originates from an HTTP POST and must not contain
+    # pickle/type markers that would reach pickle.loads().
+    response_data = strip_pickle_markers(response_data)
     if response_data is None:
         return None
 
@@ -963,9 +976,9 @@ def _deserialize_hitl_response(response_data: Any, response_type_str: str | None
         logger.debug("Response data is not a dict, returning as-is: %s", type(response_data).__name__)
         return response_data
 
-    # Try to deserialize using the type hint
+    # Try to reconstruct using the type hint (Pydantic / dataclass)
     if response_type_str:
-        response_type = _resolve_type(response_type_str)
+        response_type = resolve_type(response_type_str)
         if response_type:
             logger.debug("Found response type %s, attempting reconstruction", response_type)
             result = reconstruct_to_type(response_data, response_type)
@@ -973,6 +986,8 @@ def _deserialize_hitl_response(response_data: Any, response_type_str: str | None
             return result
         logger.warning("Could not resolve response type: %s", response_type_str)
 
-    # Fall back to generic deserialization
-    logger.debug("Falling back to generic deserialization")
-    return deserialize_value(response_data)
+    # No type hint available - return the sanitized dict as-is.
+    # We intentionally do NOT call deserialize_value() here because HITL
+    # response data is untrusted and must never flow into pickle.loads().
+    logger.debug("No type hint; returning sanitized data as-is")
+    return response_data  # type: ignore[reportUnknownVariableType]

python/packages/azurefunctions/pyproject.toml

@@ -4,7 +4,7 @@ description = "Azure Functions integration for Microsoft Agent Framework."
 authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
 readme = "README.md"
 requires-python = ">=3.10"
-version = "1.0.0b260225"
+version = "1.0.0b260521"
 license-files = ["LICENSE"]
 urls.homepage = "https://aka.ms/agent-framework"
 urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
@@ -22,10 +22,10 @@ classifiers = [
   "Typing :: Typed",
 ]
 dependencies = [
-    "agent-framework-core>=1.0.0rc2",
-    "agent-framework-durabletask",
-    "azure-functions",
-    "azure-functions-durable",
+    "agent-framework-core>=1.6.0,<2",
+    "agent-framework-durabletask>=1.0.0b260521,<2",
+    "azure-functions>=1.24.0,<2",
+    "azure-functions-durable>=1.3.1,<2",
 ]
 
 [dependency-groups]
@@ -67,6 +67,7 @@ omit = [
 
 [tool.pyright]
 extends = "../../pyproject.toml"
+include = ["agent_framework_azurefunctions"]
 
 [tool.mypy]
 plugins = ['pydantic.mypy']
@@ -90,9 +91,13 @@ exclude_dirs = ["tests"]
 executor.type = "uv"
 include = "../../shared_tasks.toml"
 
-[tool.poe.tasks]
-mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_azurefunctions"
-test = "pytest --cov=agent_framework_azurefunctions --cov-report=term-missing:skip-covered tests"
+[tool.poe.tasks.mypy]
+help = "Run MyPy for this package."
+cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_azurefunctions"
+
+[tool.poe.tasks.test]
+help = "Run the default unit test suite for this package."
+cmd = 'pytest -m "not integration" --cov=agent_framework_azurefunctions --cov-report=term-missing:skip-covered tests'
 
 [build-system]
 requires = ["flit-core >= 3.11,<4.0"]

python/packages/azurefunctions/tests/integration_tests/.env.example

@@ -1,6 +1,6 @@
 # Azure OpenAI Configuration
 AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
-AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=your-deployment-name
+AZURE_OPENAI_MODEL=your-deployment-name
 FUNCTIONS_WORKER_RUNTIME=python
 
 # Azure Functions Configuration

python/packages/azurefunctions/tests/integration_tests/README.md

@@ -14,7 +14,7 @@ cp .env.example .env
 
 Required variables:
 - `AZURE_OPENAI_ENDPOINT`
-- `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`
+- `AZURE_OPENAI_MODEL`
 - `AZURE_OPENAI_API_KEY`
 - `AzureWebJobsStorage`
 - `DURABLE_TASK_SCHEDULER_CONNECTION_STRING`

python/packages/azurefunctions/tests/integration_tests/conftest.py

@@ -111,13 +111,17 @@ def _should_skip_azure_functions_integration_tests() -> tuple[bool, str]:
             f"Durable Task Scheduler emulator not running on port {_DTS_EMULATOR_PORT}. Start with: docker run -d -p 8080:8080 -p 8082:8082 mcr.microsoft.com/dts/dts-emulator:latest",  # noqa: E501
         )
 
-    endpoint = os.getenv("AZURE_OPENAI_ENDPOINT", "").strip()
-    if not endpoint or endpoint == "https://your-resource.openai.azure.com/":
-        return True, "No real AZURE_OPENAI_ENDPOINT provided; skipping integration tests."
-
-    deployment_name = os.getenv("AZURE_OPENAI_CHAT_DEPLOYMENT_NAME", "").strip()
-    if not deployment_name or deployment_name == "your-deployment-name":
-        return True, "No real AZURE_OPENAI_CHAT_DEPLOYMENT_NAME provided; skipping integration tests."
+    has_foundry_config = bool(os.getenv("FOUNDRY_PROJECT_ENDPOINT", "").strip()) and bool(
+        os.getenv("FOUNDRY_MODEL", "").strip()
+    )
+    has_azure_openai_config = bool(os.getenv("AZURE_OPENAI_ENDPOINT", "").strip()) and bool(
+        os.getenv("AZURE_OPENAI_MODEL", "").strip()
+    )
+    if not has_foundry_config and not has_azure_openai_config:
+        return (
+            True,
+            "No real FOUNDRY_* or AZURE_OPENAI_* configuration provided; skipping integration tests.",
+        )
 
     return False, "Integration tests enabled."
 
@@ -322,22 +326,22 @@ def _is_port_in_use(port: int, host: str = _DEFAULT_HOST) -> bool:
         return sock.connect_ex((host, port)) == 0
 
 
-def _load_and_validate_env() -> None:
+def _load_and_validate_env(sample_path: Path) -> None:
     """Load .env file from current directory if it exists, then validate required environment variables.
 
     Raises pytest.fail if required environment variables are missing.
     """
     _load_env_file_if_present()
-
     # Required environment variables for Azure Functions samples
-    # These match the variables defined in .env.example
     required_env_vars = [
-        "AZURE_OPENAI_ENDPOINT",
-        "AZURE_OPENAI_CHAT_DEPLOYMENT_NAME",
         "AzureWebJobsStorage",
         "DURABLE_TASK_SCHEDULER_CONNECTION_STRING",
         "FUNCTIONS_WORKER_RUNTIME",
     ]
+    if sample_path.name == "11_workflow_parallel":
+        required_env_vars.extend(["AZURE_OPENAI_ENDPOINT", "AZURE_OPENAI_MODEL"])
+    else:
+        required_env_vars.extend(["FOUNDRY_PROJECT_ENDPOINT", "FOUNDRY_MODEL"])
 
     # Check if required env vars are set
     missing_vars = [var for var in required_env_vars if not os.environ.get(var)]
@@ -541,7 +545,7 @@ def function_app_for_test(request: pytest.FixtureRequest) -> Iterator[dict[str,
     assert sample_path is not None, "Sample path must be resolved before starting the function app"
 
     # Load .env file if it exists and validate required env vars
-    _load_and_validate_env()
+    _load_and_validate_env(sample_path)
 
     max_attempts = 3
     # The overall budget MUST be shorter than the pytest-timeout value

python/packages/azurefunctions/tests/integration_tests/test_03_reliable_streaming.py

@@ -26,7 +26,6 @@ pytestmark = [
     pytest.mark.integration,
     pytest.mark.sample("03_reliable_streaming"),
     pytest.mark.usefixtures("function_app_for_test"),
-    pytest.mark.skip(reason="Temp disabled to fix test instability - needs investigation into root cause"),
 ]
 
 
@@ -56,12 +55,11 @@ class TestSampleReliableStreaming:
         # Wait a moment for the agent to start writing to Redis
         time.sleep(2)
 
-        # Stream response from Redis with shorter timeout
-        # Note: We use text/plain to avoid SSE parsing complexity
+        # Stream response from Redis with longer timeout to account for LLM latency
         stream_response = requests.get(
             f"{self.stream_url}/{thread_id}",
             headers={"Accept": "text/plain"},
-            timeout=30,  # Shorter timeout for test
+            timeout=60,
         )
         assert stream_response.status_code == 200
 
@@ -83,7 +81,7 @@ class TestSampleReliableStreaming:
         stream_response = requests.get(
             f"{self.stream_url}/{thread_id}",
             headers={"Accept": "text/event-stream"},
-            timeout=30,  # Shorter timeout
+            timeout=60,
         )
         assert stream_response.status_code == 200
         content_type = stream_response.headers.get("content-type", "")

python/packages/azurefunctions/tests/integration_tests/test_11_workflow_parallel.py

@@ -42,6 +42,7 @@ class TestWorkflowParallel:
         self.base_url = base_url
         self.helper = sample_helper
 
+    @pytest.mark.skip(reason="xdist distributes module tests across workers, each spawning a func process")
     def test_parallel_workflow_document_analysis(self) -> None:
         """Test parallel workflow with a standard document."""
         payload = {
@@ -70,6 +71,7 @@ class TestWorkflowParallel:
         assert status["runtimeStatus"] == "Completed"
         assert "output" in status
 
+    @pytest.mark.skip(reason="xdist distributes module tests across workers, each spawning a func process")
     def test_parallel_workflow_short_document(self) -> None:
         """Test parallel workflow with a short document."""
         payload = {
@@ -89,6 +91,7 @@ class TestWorkflowParallel:
         assert status["runtimeStatus"] == "Completed"
         assert "output" in status
 
+    @pytest.mark.skip(reason="xdist distributes module tests across workers, each spawning a func process")
     def test_parallel_workflow_technical_document(self) -> None:
         """Test parallel workflow with a technical document."""
         payload = {
@@ -112,6 +115,7 @@ class TestWorkflowParallel:
         status = self.helper.wait_for_orchestration_with_output(data["statusQueryGetUri"], max_wait=300)
         assert status["runtimeStatus"] == "Completed"
 
+    @pytest.mark.skip(reason="xdist distributes module tests across workers, each spawning a func process")
     def test_workflow_status_endpoint(self) -> None:
         """Test that the workflow status endpoint works correctly."""
         payload = {

python/packages/azurefunctions/tests/test_app.py

@@ -26,6 +26,7 @@ from agent_framework_durabletask import (
 
 from agent_framework_azurefunctions import AgentFunctionApp
 from agent_framework_azurefunctions._entities import create_agent_entity
+from agent_framework_azurefunctions._workflow import SOURCE_ORCHESTRATOR
 
 FuncT = TypeVar("FuncT", bound=Callable[..., Any])
 
@@ -356,7 +357,7 @@ class TestAgentEntityOperations:
         """Test that entity can run agent operation."""
         mock_agent = Mock()
         mock_agent.run = AsyncMock(
-            return_value=AgentResponse(messages=[Message(role="assistant", text="Test response")])
+            return_value=AgentResponse(messages=[Message(role="assistant", contents=["Test response"])])
         )
 
         entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="test-conv-123"))
@@ -373,7 +374,9 @@ class TestAgentEntityOperations:
     async def test_entity_stores_conversation_history(self) -> None:
         """Test that the entity stores conversation history."""
         mock_agent = Mock()
-        mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response 1")]))
+        mock_agent.run = AsyncMock(
+            return_value=AgentResponse(messages=[Message(role="assistant", contents=["Response 1"])])
+        )
 
         entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
 
@@ -405,7 +408,9 @@ class TestAgentEntityOperations:
     async def test_entity_increments_message_count(self) -> None:
         """Test that the entity increments the message count."""
         mock_agent = Mock()
-        mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response")]))
+        mock_agent.run = AsyncMock(
+            return_value=AgentResponse(messages=[Message(role="assistant", contents=["Response"])])
+        )
 
         entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
 
@@ -444,7 +449,9 @@ class TestAgentEntityFactory:
     def test_entity_function_handles_run_operation(self) -> None:
         """Test that the entity function handles the run operation."""
         mock_agent = Mock()
-        mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response")]))
+        mock_agent.run = AsyncMock(
+            return_value=AgentResponse(messages=[Message(role="assistant", contents=["Response"])])
+        )
 
         entity_function = create_agent_entity(mock_agent)
 
@@ -469,7 +476,9 @@ class TestAgentEntityFactory:
     def test_entity_function_handles_run_agent_operation(self) -> None:
         """Test that the entity function handles the deprecated run_agent operation for backward compatibility."""
         mock_agent = Mock()
-        mock_agent.run = AsyncMock(return_value=AgentResponse(messages=[Message(role="assistant", text="Response")]))
+        mock_agent.run = AsyncMock(
+            return_value=AgentResponse(messages=[Message(role="assistant", contents=["Response"])])
+        )
 
         entity_function = create_agent_entity(mock_agent)
 
@@ -1441,5 +1450,286 @@ class TestAgentFunctionAppWorkflow:
         assert "instance-456" in url
 
 
+def _compute_state_updates(original_snapshot: dict[str, Any], current_state: dict[str, Any]) -> dict[str, Any]:
+    """Compute state updates by comparing current state against the original snapshot.
+
+    This mirrors the inlined logic in ``_app.py``'s ``executor_activity.run()``.
+    """
+    original_keys = set(original_snapshot.keys())
+    current_keys = set(current_state.keys())
+    updates: dict[str, Any] = {}
+    for key in current_keys:
+        if key not in original_keys or current_state[key] != original_snapshot.get(key):
+            updates[key] = current_state[key]
+    return updates
+
+
+class TestStateSnapshotDiff:
+    """Test suite for state snapshot diffing in activity execution.
+
+    The activity executor snapshots state before execution and diffs against the
+    post-execution state to determine which keys were updated. These tests exercise
+    the production snapshot helper and the state-update diffing logic to ensure that
+    in-place mutations to nested objects (dicts, lists) are correctly detected as changes.
+    """
+
+    def test_nested_dict_mutation_detected_in_diff(self) -> None:
+        """Test that mutating values inside a nested dict appears in the diff."""
+        from agent_framework._workflows._state import State
+
+        from agent_framework_azurefunctions._app import _create_state_snapshot
+
+        deserialized_state: dict[str, Any] = {
+            "Local.config": {"code": "", "enabled": False},
+            "simple_key": "simple_value",
+        }
+
+        original_snapshot = _create_state_snapshot(deserialized_state)
+
+        shared_state = State()
+        shared_state.import_state(deserialized_state)
+
+        config = shared_state.get("Local.config")
+        config["code"] = "SOMECODEXXX"
+        config["enabled"] = True
+
+        shared_state.commit()
+        current_state = shared_state.export_state()
+
+        updates = _compute_state_updates(original_snapshot, current_state)
+
+        assert "Local.config" in updates
+        assert updates["Local.config"]["code"] == "SOMECODEXXX"
+        assert updates["Local.config"]["enabled"] is True
+
+    def test_new_key_in_nested_dict_detected_in_diff(self) -> None:
+        """Test that adding a key to a nested dict appears in the diff."""
+        from agent_framework._workflows._state import State
+
+        from agent_framework_azurefunctions._app import _create_state_snapshot
+
+        deserialized_state: dict[str, Any] = {
+            "Local.data": {"existing": "value"},
+        }
+
+        original_snapshot = _create_state_snapshot(deserialized_state)
+
+        shared_state = State()
+        shared_state.import_state(deserialized_state)
+
+        data = shared_state.get("Local.data")
+        data["code"] = "NEW_CODE"
+
+        shared_state.commit()
+        current_state = shared_state.export_state()
+
+        updates = _compute_state_updates(original_snapshot, current_state)
+
+        assert "Local.data" in updates
+        assert updates["Local.data"]["code"] == "NEW_CODE"
+
+    def test_nested_list_mutation_detected_in_diff(self) -> None:
+        """Test that appending to a nested list appears in the diff."""
+        from agent_framework._workflows._state import State
+
+        from agent_framework_azurefunctions._app import _create_state_snapshot
+
+        deserialized_state: dict[str, Any] = {
+            "Local.items": [1, 2, 3],
+        }
+
+        original_snapshot = _create_state_snapshot(deserialized_state)
+
+        shared_state = State()
+        shared_state.import_state(deserialized_state)
+
+        items = shared_state.get("Local.items")
+        items.append(4)
+
+        shared_state.commit()
+        current_state = shared_state.export_state()
+
+        updates = _compute_state_updates(original_snapshot, current_state)
+
+        assert "Local.items" in updates
+        assert updates["Local.items"] == [1, 2, 3, 4]
+
+    def test_new_top_level_key_detected_in_diff(self) -> None:
+        """Test that setting a new top-level key appears in the diff."""
+        from agent_framework._workflows._state import State
+
+        from agent_framework_azurefunctions._app import _create_state_snapshot
+
+        deserialized_state: dict[str, Any] = {
+            "existing": "value",
+        }
+
+        original_snapshot = _create_state_snapshot(deserialized_state)
+
+        shared_state = State()
+        shared_state.import_state(deserialized_state)
+
+        shared_state.set("Local.code", "SOMECODEXXX")
+
+        shared_state.commit()
+        current_state = shared_state.export_state()
+
+        updates = _compute_state_updates(original_snapshot, current_state)
+
+        assert "Local.code" in updates
+        assert updates["Local.code"] == "SOMECODEXXX"
+
+    def test_unchanged_nested_state_produces_empty_diff(self) -> None:
+        """Test that unmodified nested state produces no updates."""
+        from agent_framework._workflows._state import State
+
+        from agent_framework_azurefunctions._app import _create_state_snapshot
+
+        deserialized_state: dict[str, Any] = {
+            "Local.config": {"code": "existing", "enabled": True},
+            "simple_key": "simple_value",
+        }
+
+        original_snapshot = _create_state_snapshot(deserialized_state)
+
+        shared_state = State()
+        shared_state.import_state(deserialized_state)
+
+        # No mutations performed
+        shared_state.commit()
+        current_state = shared_state.export_state()
+
+        updates = _compute_state_updates(original_snapshot, current_state)
+
+        assert updates == {}
+
+    def test_shallow_copy_would_miss_nested_mutations(self) -> None:
+        """Regression test: a shallow copy (dict()) shares nested refs, hiding mutations.
+
+        This reproduces the original bug from #4500 where ``dict(deserialized_state)``
+        was used instead of ``copy.deepcopy()``. With a shallow copy the snapshot and
+        the live state share nested objects, so in-place mutations appear in both and
+        the diff produces an empty update set.
+        """
+        from agent_framework._workflows._state import State
+
+        deserialized_state: dict[str, Any] = {
+            "Local.config": {"code": "", "enabled": False},
+        }
+
+        # Shallow copy (the OLD, buggy behaviour)
+        shallow_snapshot = dict(deserialized_state)
+
+        shared_state = State()
+        shared_state.import_state(deserialized_state)
+
+        config = shared_state.get("Local.config")
+        config["code"] = "SOMECODEXXX"
+        config["enabled"] = True
+
+        shared_state.commit()
+        current_state = shared_state.export_state()
+
+        # With a shallow copy the mutation leaks into the snapshot → empty diff
+        updates_shallow = _compute_state_updates(shallow_snapshot, current_state)
+        assert updates_shallow == {}, "shallow copy should miss nested mutations (demonstrating the bug)"
+
+    def test_create_state_snapshot_isolates_nested_objects(self) -> None:
+        """Verify _create_state_snapshot produces a deep copy that is mutation-proof.
+
+        This ensures the production snapshot helper is not equivalent to ``dict()``
+        and will correctly isolate nested objects so that later mutations are detected.
+        """
+        from agent_framework_azurefunctions._app import _create_state_snapshot
+
+        original: dict[str, Any] = {
+            "nested_dict": {"a": 1},
+            "nested_list": [1, 2, 3],
+        }
+
+        snapshot = _create_state_snapshot(original)
+
+        # Mutate the originals in place
+        original["nested_dict"]["a"] = 999
+        original["nested_list"].append(4)
+
+        # Snapshot must be unaffected
+        assert snapshot["nested_dict"]["a"] == 1
+        assert snapshot["nested_list"] == [1, 2, 3]
+
+    def test_executor_activity_detects_nested_state_mutations(self) -> None:
+        """Integration test: the full activity wrapper detects nested mutations.
+
+        This exercises the actual executor_activity function registered by
+        _setup_executor_activity to verify the production code path uses
+        _create_state_snapshot (deep copy) rather than dict() (shallow copy).
+        If the implementation regressed to using a shallow copy such as
+        ``dict(deserialized_state)``, this test would fail because in-place
+        mutations would leak into the snapshot and produce an empty diff.
+        """
+        mock_executor = Mock()
+        mock_executor.id = "test-exec"
+
+        async def mutate_nested_state(
+            message: Any,
+            source_executor_ids: Any,
+            state: Any,
+            runner_context: Any,
+        ) -> None:
+            config = state.get("Local.config")
+            config["code"] = "MUTATED"
+            config["enabled"] = True
+            state.commit()
+
+        mock_executor.execute = AsyncMock(side_effect=mutate_nested_state)
+
+        mock_workflow = Mock()
+        mock_workflow.executors = {"test-exec": mock_executor}
+
+        # Capture the activity function by making decorators pass-through
+        captured_activity: dict[str, Any] = {}
+
+        def passthrough_function_name(name: str) -> Callable[[FuncT], FuncT]:
+            def decorator(fn: FuncT) -> FuncT:
+                captured_activity["fn"] = fn
+                return fn
+
+            return decorator
+
+        def passthrough_activity_trigger(input_name: str) -> Callable[[FuncT], FuncT]:
+            def decorator(fn: FuncT) -> FuncT:
+                return fn
+
+            return decorator
+
+        with (
+            patch.object(AgentFunctionApp, "function_name", side_effect=passthrough_function_name),
+            patch.object(AgentFunctionApp, "activity_trigger", side_effect=passthrough_activity_trigger),
+            patch.object(AgentFunctionApp, "_setup_workflow_orchestration"),
+        ):
+            AgentFunctionApp(workflow=mock_workflow)
+
+        assert "fn" in captured_activity, "activity function was not captured"
+
+        # Call the activity with nested state that the executor will mutate
+        input_data = json.dumps({
+            "message": "test",
+            "shared_state_snapshot": {
+                "Local.config": {"code": "", "enabled": False},
+            },
+            "source_executor_ids": [SOURCE_ORCHESTRATOR],
+        })
+
+        result = json.loads(captured_activity["fn"](input_data))
+
+        # The deep copy snapshot must detect the in-place nested mutations
+        assert "Local.config" in result["shared_state_updates"], (
+            "nested mutation not detected — snapshot may be using shallow copy"
+        )
+        updated_config = result["shared_state_updates"]["Local.config"]
+        assert updated_config["code"] == "MUTATED"
+        assert updated_config["enabled"] is True
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "--tb=short"])

python/packages/azurefunctions/tests/test_entities.py

@@ -19,7 +19,9 @@ FuncT = TypeVar("FuncT", bound=Callable[..., Any])
 
 def _agent_response(text: str | None) -> AgentResponse:
     """Create an AgentResponse with a single assistant message."""
-    message = Message(role="assistant", text=text) if text is not None else Message(role="assistant", text="")
+    message = (
+        Message(role="assistant", contents=[text]) if text is not None else Message(role="assistant", contents=[""])
+    )
     return AgentResponse(messages=[message])
 
 

python/packages/azurefunctions/tests/test_func_utils.py

@@ -21,6 +21,7 @@ from agent_framework_azurefunctions._serialization import (
     deserialize_value,
     reconstruct_to_type,
     serialize_value,
+    strip_pickle_markers,
 )
 
 
@@ -106,7 +107,7 @@ class TestCapturingRunnerContext:
     @pytest.mark.asyncio
     async def test_add_event_queues_event(self, context: CapturingRunnerContext) -> None:
         """Test that add_event queues events correctly."""
-        event = WorkflowEvent.output(executor_id="exec_1", data="output")
+        event = WorkflowEvent("output", executor_id="exec_1", data="output")
 
         await context.add_event(event)
 
@@ -119,7 +120,7 @@ class TestCapturingRunnerContext:
     @pytest.mark.asyncio
     async def test_drain_events_clears_queue(self, context: CapturingRunnerContext) -> None:
         """Test that drain_events clears the event queue."""
-        await context.add_event(WorkflowEvent.output(executor_id="e", data="test"))
+        await context.add_event(WorkflowEvent("output", executor_id="e", data="test"))
 
         await context.drain_events()  # First drain
         events = await context.drain_events()  # Second drain
@@ -131,14 +132,14 @@ class TestCapturingRunnerContext:
         """Test has_events returns correct boolean."""
         assert await context.has_events() is False
 
-        await context.add_event(WorkflowEvent.output(executor_id="e", data="test"))
+        await context.add_event(WorkflowEvent("output", executor_id="e", data="test"))
 
         assert await context.has_events() is True
 
     @pytest.mark.asyncio
     async def test_next_event_waits_for_event(self, context: CapturingRunnerContext) -> None:
         """Test that next_event returns queued events."""
-        event = WorkflowEvent.output(executor_id="e", data="waited")
+        event = WorkflowEvent("output", executor_id="e", data="waited")
         await context.add_event(event)
 
         result = await context.next_event()
@@ -170,7 +171,7 @@ class TestCapturingRunnerContext:
     async def test_reset_for_new_run_clears_state(self, context: CapturingRunnerContext) -> None:
         """Test that reset_for_new_run clears all state."""
         await context.send_message(WorkflowMessage(data="test", target_id="t", source_id="s"))
-        await context.add_event(WorkflowEvent.output(executor_id="e", data="event"))
+        await context.add_event(WorkflowEvent("output", executor_id="e", data="event"))
         context.set_streaming(True)
 
         context.reset_for_new_run()
@@ -205,7 +206,7 @@ class TestSerializationRoundtrip:
 
     def test_roundtrip_chat_message(self) -> None:
         """Test Message survives encode → decode roundtrip."""
-        original = Message(role="user", text="Hello")
+        original = Message(role="user", contents=["Hello"])
         encoded = serialize_value(original)
         decoded = deserialize_value(encoded)
 
@@ -215,7 +216,7 @@ class TestSerializationRoundtrip:
     def test_roundtrip_agent_executor_request(self) -> None:
         """Test AgentExecutorRequest with nested Messages roundtrips."""
         original = AgentExecutorRequest(
-            messages=[Message(role="user", text="Hi")],
+            messages=[Message(role="user", contents=["Hi"])],
             should_respond=True,
         )
         encoded = serialize_value(original)
@@ -230,7 +231,8 @@ class TestSerializationRoundtrip:
         """Test AgentExecutorResponse with nested AgentResponse roundtrips."""
         original = AgentExecutorResponse(
             executor_id="test_exec",
-            agent_response=AgentResponse(messages=[Message(role="assistant", text="Reply")]),
+            agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Reply"])]),
+            full_conversation=[Message(role="assistant", contents=["Reply"])],
         )
         encoded = serialize_value(original)
         decoded = deserialize_value(encoded)
@@ -270,8 +272,8 @@ class TestSerializationRoundtrip:
     def test_roundtrip_list_of_objects(self) -> None:
         """Test list of typed objects roundtrips."""
         original = [
-            Message(role="user", text="Q"),
-            Message(role="assistant", text="A"),
+            Message(role="user", contents=["Q"]),
+            Message(role="assistant", contents=["A"]),
         ]
         encoded = serialize_value(original)
         decoded = deserialize_value(encoded)
@@ -282,7 +284,7 @@ class TestSerializationRoundtrip:
 
     def test_roundtrip_dict_of_objects(self) -> None:
         """Test dict with typed values roundtrips (used for shared state)."""
-        original = {"count": 42, "msg": Message(role="user", text="Hi")}
+        original = {"count": 42, "msg": Message(role="user", contents=["Hi"])}
         encoded = serialize_value(original)
         decoded = deserialize_value(encoded)
 
@@ -353,7 +355,11 @@ class TestReconstructToType:
         assert result.comment == "Great"
 
     def test_reconstruct_from_checkpoint_markers(self) -> None:
-        """Test that data with checkpoint markers is decoded via deserialize_value."""
+        """Test that data with checkpoint markers is decoded via deserialize_value.
+
+        reconstruct_to_type is general-purpose and handles trusted checkpoint
+        data.  Untrusted HITL callers must call strip_pickle_markers() first.
+        """
         original = SampleData(value=99, name="marker-test")
         encoded = serialize_value(original)
 
@@ -372,3 +378,73 @@ class TestReconstructToType:
         result = reconstruct_to_type(data, Unrelated)
 
         assert result == data
+
+    def test_reconstruct_strips_injected_pickle_markers(self) -> None:
+        """End-to-end: strip_pickle_markers + reconstruct_to_type blocks attack.
+
+        This mirrors the real HITL flow where callers sanitize before reconstruction.
+        """
+        malicious = {"__pickled__": "gASVDgAAAAAAAACMBHRlc3SULg==", "__type__": "builtins:str"}
+        sanitized = strip_pickle_markers(malicious)
+        result = reconstruct_to_type(sanitized, str)
+        assert result is None
+
+
+class TestStripPickleMarkers:
+    """Security tests for strip_pickle_markers — the defence-in-depth layer
+    that prevents untrusted HTTP input from reaching pickle.loads()."""
+
+    def test_strips_top_level_pickle_marker(self) -> None:
+        """A dict containing __pickled__ must be replaced with None."""
+        data = {"__pickled__": "PAYLOAD", "__type__": "os:system"}
+        assert strip_pickle_markers(data) is None
+
+    def test_strips_top_level_type_marker_only(self) -> None:
+        """Even __type__ alone (without __pickled__) must be neutralised."""
+        data = {"__type__": "os:system", "other": "value"}
+        assert strip_pickle_markers(data) is None
+
+    def test_strips_nested_pickle_marker(self) -> None:
+        """Pickle markers nested inside a dict must be neutralised."""
+        data = {"safe": "value", "nested": {"__pickled__": "PAYLOAD", "__type__": "os:system"}}
+        result = strip_pickle_markers(data)
+        assert result == {"safe": "value", "nested": None}
+
+    def test_strips_pickle_marker_in_list(self) -> None:
+        """Pickle markers inside a list element must be neutralised."""
+        data = [{"__pickled__": "PAYLOAD"}, "safe"]
+        result = strip_pickle_markers(data)
+        assert result == [None, "safe"]
+
+    def test_strips_deeply_nested_marker(self) -> None:
+        """Deeply nested pickle markers must be neutralised."""
+        data = {"a": {"b": {"c": {"__pickled__": "deep"}}}}
+        result = strip_pickle_markers(data)
+        assert result == {"a": {"b": {"c": None}}}
+
+    def test_preserves_safe_dict(self) -> None:
+        """Dicts without pickle markers must be left untouched."""
+        data = {"approved": True, "reason": "Looks good"}
+        assert strip_pickle_markers(data) == data
+
+    def test_preserves_primitives(self) -> None:
+        """Primitive values must pass through unchanged."""
+        assert strip_pickle_markers("hello") == "hello"
+        assert strip_pickle_markers(42) == 42
+        assert strip_pickle_markers(None) is None
+        assert strip_pickle_markers(True) is True
+
+    def test_preserves_safe_list(self) -> None:
+        """Lists without pickle markers must be left untouched."""
+        data = [1, "two", {"key": "value"}]
+        assert strip_pickle_markers(data) == data
+
+    def test_mixed_safe_and_malicious(self) -> None:
+        """Only the malicious entries should be stripped; safe entries remain."""
+        data = {
+            "user_input": "hello",
+            "evil": {"__pickled__": "PAYLOAD", "__type__": "os:system"},
+            "count": 42,
+        }
+        result = strip_pickle_markers(data)
+        assert result == {"user_input": "hello", "evil": None, "count": 42}

python/packages/azurefunctions/tests/test_orchestration.py

@@ -155,7 +155,7 @@ class TestAgentResponseHelpers:
 
         # Simulate successful entity task completion
         entity_task.state = TaskState.SUCCEEDED
-        entity_task.result = AgentResponse(messages=[Message(role="assistant", text="Test response")]).to_dict()
+        entity_task.result = AgentResponse(messages=[Message(role="assistant", contents=["Test response"])]).to_dict()
 
         # Clear pending_tasks to simulate that parent has processed the child
         task.pending_tasks.clear()
@@ -197,7 +197,9 @@ class TestAgentResponseHelpers:
 
         # Simulate successful entity task with JSON response
         entity_task.state = TaskState.SUCCEEDED
-        entity_task.result = AgentResponse(messages=[Message(role="assistant", text='{"answer": "42"}')]).to_dict()
+        entity_task.result = AgentResponse(
+            messages=[Message(role="assistant", contents=['{"answer": "42"}'])]
+        ).to_dict()
 
         # Clear pending_tasks to simulate that parent has processed the child
         task.pending_tasks.clear()

python/packages/azurefunctions/tests/test_workflow.py

@@ -177,10 +177,10 @@ class TestBuildAgentExecutorResponse:
         # Create a previous response with conversation history
         previous = AgentExecutorResponse(
             executor_id="prev",
-            agent_response=AgentResponse(messages=[Message(role="assistant", text="Previous")]),
+            agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Previous"])]),
             full_conversation=[
-                Message(role="user", text="First"),
-                Message(role="assistant", text="Previous"),
+                Message(role="user", contents=["First"]),
+                Message(role="assistant", contents=["Previous"]),
             ],
         )
 
@@ -211,7 +211,8 @@ class TestExtractMessageContent:
         """Test extracting from AgentExecutorResponse with text."""
         response = AgentExecutorResponse(
             executor_id="exec",
-            agent_response=AgentResponse(messages=[Message(role="assistant", text="Response text")]),
+            agent_response=AgentResponse(messages=[Message(role="assistant", contents=["Response text"])]),
+            full_conversation=[Message(role="assistant", contents=["Response text"])],
         )
 
         result = _extract_message_content(response)
@@ -224,10 +225,14 @@ class TestExtractMessageContent:
             executor_id="exec",
             agent_response=AgentResponse(
                 messages=[
-                    Message(role="user", text="First"),
-                    Message(role="assistant", text="Last message"),
+                    Message(role="user", contents=["First"]),
+                    Message(role="assistant", contents=["Last message"]),
                 ]
             ),
+            full_conversation=[
+                Message(role="user", contents=["First"]),
+                Message(role="assistant", contents=["Last message"]),
+            ],
         )
 
         result = _extract_message_content(response)
@@ -239,8 +244,8 @@ class TestExtractMessageContent:
         """Test extracting from AgentExecutorRequest."""
         request = AgentExecutorRequest(
             messages=[
-                Message(role="user", text="First"),
-                Message(role="user", text="Last request"),
+                Message(role="user", contents=["First"]),
+                Message(role="user", contents=["Last request"]),
             ]
         )