.NET: Python: Azure Functions feature branch (#1916)

* Python: Add Scaffolding for Durable AzureFunctions package to Agent Framework (#1823) * Add scafolding * update readme * add code owners and label * update owners * .NET: Durable extension: initial src and unit tests (#1900) * Python: Add Durable Agent Wrapper code (#1913) * add initial changes * Move code and add single sample * Update logger * Remove unused code * address PR comments * cleanup code and address comments --------- Co-authored-by: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> * Azure Functions .NET samples (#1939) * Python: Add Unit tests for Azurefunctions package (#1976) * Add Unit tests for Azurefunctions * remove duplicate import * .NET: [Feature Branch] Migrate state schema updates and support for agents as MCP tools (#1979) * Python: Add more samples for Azure Functions (#1980) * Move all samples * fix comments * remove dead lines * Make samples simpler * .NET: [Feature Branch] Durable Task extension integration tests (#2017) * .NET: [Feature Branch] Update OpenAI config for integration tests (#2063) * Python: Add Integration tests for AzureFunctions (#2020) * Add Integration tests * Remove DTS extension * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Add pyi file for type safety * Add samples in readme * Updated all readme instructions * Address comments * Update readmes * Fix requirements * Address comments --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * .NET: [Feature Branch] Update dotnet-build-and-test.yml to support integration tests (#2070) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Fix DTS startup issue and improve logging (#2103) * .NET: [Feature Branch] Introduce Azure OpenAI config for .NET pipeline (#2106) Also fixes an issue where we were trying to start docker containers for integration tests on Windows, which doesn't work. Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Fix uv.lock after merge * Python: Add README for Azure Functions samples setup (#2100) * Add README for Azure Functions samples setup Added setup instructions for Azure Functions samples, including environment setup, virtual environment creation, and running samples. * Update python/samples/getting_started/azure_functions/README.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Laveesh Rohra <larohra@microsoft.com> * Fix or remove broken markdown file links (#2115) * .NET: [Feature Branch] Update HTTP API to be consistent across languages (#2118) * Python: Fix AzureFunctions Integration Tests (#2116) * Add Identity Auth to samples * Update python/samples/getting_started/azure_functions/README.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/samples/getting_started/azure_functions/01_single_agent/function_app.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/samples/getting_started/azure_functions/02_multi_agent/function_app.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/samples/getting_started/azure_functions/06_multi_agent_orchestration_conditionals/README.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Python: Fix Http Schema (#2112) * Rename to threadid * Respond in plain text * Make snake-case * Add http prefix * rename to wait-for-response * Add query param check * address comments * .NET: Remove IsPackable=false in preparation for nuget release (#2142) * Python: Move `azurefunctions` to `azure` for import (#2141) * Move import to Azure * fix mypy * Update python/packages/azurefunctions/README.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Add missing types * Address comments --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/azurefunctions/pyproject.toml Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/azurefunctions/agent_framework_azurefunctions/__init__.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Fix imports * Address PR feedback from westey-m (#2150) - Adds a link from the /dotnet/samples/README.md to /dotnet/samples/AzureFunctions - Make DurableAgentThread deserialization internal for future-proofing - Update JSON serialization logic to address recently discovered issues with source generator serialization * Address comments (#2160) --------- Co-authored-by: Laveesh Rohra <larohra@microsoft.com> Co-authored-by: Chris Gillum <cgillum@microsoft.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Anirudh Garg <anirudhg@microsoft.com>
2026-06-16 21:04:09 +08:00 · 2025-11-13 02:00:53 +00:00
parent 5537b1da79
commit 67a8147151
216 changed files with 18178 additions and 2 deletions
@@ -0,0 +1,11 @@
+# Azure OpenAI Configuration
+AZURE_OPENAI_ENDPOINT=https://your-resource.openai.azure.com/
+AZURE_OPENAI_CHAT_DEPLOYMENT_NAME=your-deployment-name
+FUNCTIONS_WORKER_RUNTIME=python
+RUN_INTEGRATION_TESTS=true
+
+# Azure Functions Configuration
+AzureWebJobsStorage=UseDevelopmentStorage=true
+DURABLE_TASK_SCHEDULER_CONNECTION_STRING=Endpoint=http://localhost:8080;Authentication=None
+
+# Note: TASKHUB_NAME is not required for integration tests; it is auto-generated per test run.
@@ -0,0 +1,81 @@
+# Sample Integration Tests
+
+Integration tests that validate the Durable Agent Framework samples by running them as Azure Functions.
+
+## Setup
+
+### 1. Create `.env` file
+
+Copy `.env.example` to `.env` and fill in your Azure credentials:
+
+```bash
+cp .env.example .env
+```
+
+Required variables:
+- `AZURE_OPENAI_ENDPOINT`
+- `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`
+- `AZURE_OPENAI_API_KEY`
+- `AzureWebJobsStorage`
+- `DURABLE_TASK_SCHEDULER_CONNECTION_STRING`
+- `FUNCTIONS_WORKER_RUNTIME`
+
+### 2. Start required services
+
+**Azurite (for orchestration tests):**
+```bash
+docker run -d -p 10000:10000 -p 10001:10001 -p 10002:10002 mcr.microsoft.com/azure-storage/azurite
+```
+
+**Durable Task Scheduler:**
+```bash
+docker run -d -p 8080:8080 -p 8082:8082 mcr.microsoft.com/dts/dts-emulator:latest
+```
+
+## Running Tests
+
+The tests automatically start and stop the Azure Functions app for each sample.
+
+### Run all sample tests
+```bash
+uv run pytest packages/azurefunctions/tests/integration_tests -v
+```
+
+### Run specific sample
+```bash
+uv run pytest packages/azurefunctions/tests/integration_tests/test_01_single_agent.py -v
+```
+
+### Run with verbose output
+```bash
+uv run pytest packages/azurefunctions/tests/integration_tests -sv
+```
+
+## How It Works
+
+Each test file uses pytest markers to automatically configure and start the function app:
+
+```python
+pytestmark = [
+    pytest.mark.sample("01_single_agent"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+```
+
+The `function_app_for_test` fixture:
+1. Loads environment variables from `.env`
+2. Validates required variables are present
+3. Starts the function app on a dynamically allocated port
+4. Waits for the app to be ready
+5. Runs your tests
+6. Tears down the function app
+
+## Troubleshooting
+
+
+**Missing environment variables:**
+Ensure your `.env` file contains all required variables from `.env.example`.
+
+**Tests timeout:**
+Check that Azure OpenAI credentials are valid and the service is accessible.
@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.
@@ -0,0 +1,121 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Pytest configuration for Durable Agent Framework tests.
+
+This module provides fixtures and configuration for pytest.
+"""
+
+import subprocess
+from collections.abc import Iterator, Mapping
+from typing import Any
+
+import pytest
+import requests
+
+from .testutils import (
+    FunctionAppStartupError,
+    build_base_url,
+    cleanup_function_app,
+    find_available_port,
+    get_sample_path_from_marker,
+    load_and_validate_env,
+    start_function_app,
+    wait_for_function_app_ready,
+)
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    """Register custom markers."""
+    config.addinivalue_line("markers", "orchestration: marks tests that use orchestrations (require Azurite)")
+    config.addinivalue_line(
+        "markers",
+        "sample(path): specify the sample directory path for the test (e.g., @pytest.mark.sample('01_single_agent'))",
+    )
+
+
+@pytest.fixture(scope="session")
+def function_app_running() -> bool:
+    """
+    Check if the function app is running on localhost:7071.
+
+    This fixture can be used to skip tests if the function app is not available.
+    """
+    try:
+        response = requests.get("http://localhost:7071/api/health", timeout=2)
+        return response.status_code == 200
+    except requests.exceptions.RequestException:
+        return False
+
+
+@pytest.fixture(scope="session")
+def skip_if_no_function_app(function_app_running: bool) -> None:
+    """Skip test if function app is not running."""
+    if not function_app_running:
+        pytest.skip("Function app is not running on http://localhost:7071")
+
+
+@pytest.fixture(scope="module")
+def function_app_for_test(request: pytest.FixtureRequest) -> Iterator[dict[str, int | str]]:
+    """
+    Start the function app for the corresponding sample based on marker.
+
+    This fixture:
+    1. Determines which sample to run from @pytest.mark.sample()
+    2. Validates environment variables
+    3. Starts the function app using 'func start'
+    4. Waits for the app to be ready
+    5. Tears down the app after tests complete
+
+    Usage:
+    @pytest.mark.sample("01_single_agent")
+    @pytest.mark.usefixtures("function_app_for_test")
+    class TestSample01SingleAgent:
+        ...
+    """
+    # Get sample path from marker
+    sample_path, error_message = get_sample_path_from_marker(request)
+    if error_message:
+        pytest.fail(error_message)
+
+    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()
+
+    max_attempts = 3
+    last_error: Exception | None = None
+    func_process: subprocess.Popen[Any] | None = None
+    base_url = ""
+    port = 0
+
+    for _ in range(max_attempts):
+        port = find_available_port()
+        base_url = build_base_url(port)
+        func_process = start_function_app(sample_path, port)
+
+        try:
+            wait_for_function_app_ready(func_process, port)
+            last_error = None
+            break
+        except FunctionAppStartupError as exc:
+            last_error = exc
+            cleanup_function_app(func_process)
+            func_process = None
+
+    if func_process is None:
+        error_message = f"Function app failed to start after {max_attempts} attempt(s)."
+        if last_error is not None:
+            error_message += f" Last error: {last_error}"
+        pytest.fail(error_message)
+
+    try:
+        yield {"base_url": base_url, "port": port}
+    finally:
+        if func_process is not None:
+            cleanup_function_app(func_process)
+
+
+@pytest.fixture(scope="module")
+def base_url(function_app_for_test: Mapping[str, int | str]) -> str:
+    """Expose the function app's base URL to tests."""
+    return str(function_app_for_test["base_url"])
@@ -0,0 +1,114 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for Single Agent Sample
+
+Tests the single agent sample with various message formats and session management.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite or Azure Storage account configured
+
+Usage:
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_01_single_agent.py -v
+"""
+
+import pytest
+
+from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.sample("01_single_agent"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+class TestSampleSingleAgent:
+    """Tests for 01_single_agent sample."""
+
+    @pytest.fixture(autouse=True)
+    def _set_base_url(self, base_url: str) -> None:
+        """Provide agent-specific base URL for the tests."""
+        self.base_url = f"{base_url}/api/agents/Joker"
+
+    def test_health_check(self, base_url: str) -> None:
+        """Test health check endpoint."""
+        response = SampleTestHelper.get(f"{base_url}/api/health")
+        assert response.status_code == 200
+        data = response.json()
+        assert data["status"] == "healthy"
+
+    def test_simple_message_json(self) -> None:
+        """Test sending a simple message with JSON payload."""
+        response = SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Tell me a short joke about cloud computing.", "thread_id": "test-simple-json"},
+        )
+        # Agent can return 200 (immediate) or 202 (async with wait_for_response=false)
+        assert response.status_code in [200, 202]
+        data = response.json()
+
+        if response.status_code == 200:
+            # Synchronous response - check result directly
+            assert data["status"] == "success"
+            assert "response" in data
+            assert data["message_count"] >= 1
+        else:
+            # Async response - check we got correlation info
+            assert "correlation_id" in data or "thread_id" in data
+
+    def test_simple_message_plain_text(self) -> None:
+        """Test sending a message with plain text payload."""
+        response = SampleTestHelper.post_text(f"{self.base_url}/run", "Tell me a short joke about networking.")
+        assert response.status_code in [200, 202]
+
+        # Agent responded with plain text when the request body was text/plain.
+        assert response.text.strip()
+        assert response.headers.get("x-ms-thread-id") is not None
+
+    def test_thread_id_in_query(self) -> None:
+        """Test using thread_id in query parameter."""
+        response = SampleTestHelper.post_text(
+            f"{self.base_url}/run?thread_id=test-query-thread", "Tell me a short joke about weather in Texas."
+        )
+        assert response.status_code in [200, 202]
+
+        assert response.text.strip()
+        assert response.headers.get("x-ms-thread-id") == "test-query-thread"
+
+    def test_conversation_continuity(self) -> None:
+        """Test conversation context is maintained across requests."""
+        thread_id = "test-continuity"
+
+        # First message
+        response1 = SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Tell me a short joke about weather in Seattle.", "thread_id": thread_id},
+        )
+        assert response1.status_code in [200, 202]
+
+        if response1.status_code == 200:
+            data1 = response1.json()
+            assert data1["message_count"] == 1
+
+            # Second message in same session
+            response2 = SampleTestHelper.post_json(
+                f"{self.base_url}/run", {"message": "What about San Francisco?", "thread_id": thread_id}
+            )
+            assert response2.status_code == 200
+            data2 = response2.json()
+            assert data2["message_count"] == 2
+        else:
+            # In async mode, we can't easily test message count
+            # Just verify we can make multiple calls
+            response2 = SampleTestHelper.post_json(
+                f"{self.base_url}/run", {"message": "What about Texas?", "thread_id": thread_id}
+            )
+            assert response2.status_code == 202
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for Multi-Agent Sample
+
+Tests the multi-agent sample with different agent endpoints.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite or Azure Storage account configured
+
+Usage:
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_02_multi_agent.py -v
+"""
+
+import pytest
+
+from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.sample("02_multi_agent"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+class TestSampleMultiAgent:
+    """Tests for 02_multi_agent sample."""
+
+    @pytest.fixture(autouse=True)
+    def _set_agent_urls(self, base_url: str) -> None:
+        """Configure base URLs for Weather and Math agents."""
+        self.weather_base_url = f"{base_url}/api/agents/WeatherAgent"
+        self.math_base_url = f"{base_url}/api/agents/MathAgent"
+
+    def test_weather_agent(self) -> None:
+        """Test WeatherAgent endpoint."""
+        response = SampleTestHelper.post_json(
+            f"{self.weather_base_url}/run",
+            {"message": "What is the weather in Seattle?"},
+        )
+        assert response.status_code == 200
+        data = response.json()
+        assert data["status"] == "success"
+        assert "response" in data
+
+    def test_math_agent(self) -> None:
+        """Test MathAgent endpoint."""
+        response = SampleTestHelper.post_json(
+            f"{self.math_base_url}/run",
+            {"message": "Calculate a 20% tip on a $50 bill", "wait_for_response": False},
+        )
+        assert response.status_code == 202
+        data = response.json()
+
+        assert data["status"] == "accepted"
+        assert "correlation_id" in data
+        assert "thread_id" in data
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for Callbacks Sample
+
+Tests the callbacks sample for event tracking and management.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite or Azure Storage account configured
+
+Usage:
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_03_callbacks.py -v
+"""
+
+from typing import Any
+
+import pytest
+import requests
+
+from .testutils import (
+    TIMEOUT,
+    SampleTestHelper,
+    skip_if_azure_functions_integration_tests_disabled,
+)
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.sample("03_callbacks"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+class TestSampleCallbacks:
+    """Tests for 03_callbacks sample."""
+
+    @pytest.fixture(autouse=True)
+    def _set_base_url(self, base_url: str) -> None:
+        """Provide the callback agent base URL for each test."""
+        self.base_url = f"{base_url}/api/agents/CallbackAgent"
+
+    @staticmethod
+    def _wait_for_callback_events(base_url: str, thread_id: str) -> list[dict[str, Any]]:
+        events: list[dict[str, Any]] = []
+        response = SampleTestHelper.get(f"{base_url}/callbacks/{thread_id}")
+        if response.status_code == 200:
+            events = response.json()
+        return events
+
+    def test_agent_with_callbacks(self) -> None:
+        """Test agent execution with callback tracking."""
+        thread_id = "test-callback"
+
+        response = SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Tell me about Python", "thread_id": thread_id},
+        )
+        assert response.status_code == 200
+        data = response.json()
+
+        assert data["status"] == "success"
+
+        events = self._wait_for_callback_events(self.base_url, thread_id)
+
+        assert events
+        assert any(event.get("event_type") == "final" for event in events)
+
+    def test_get_callbacks(self) -> None:
+        """Test retrieving callback events."""
+        thread_id = "test-callback-retrieve"
+
+        # Send a message first
+        SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Hello", "thread_id": thread_id, "wait_for_response": False},
+        )
+
+        # Get callbacks
+        response = SampleTestHelper.get(f"{self.base_url}/callbacks/{thread_id}")
+        assert response.status_code == 200
+        data = response.json()
+        assert isinstance(data, list)
+
+    def test_delete_callbacks(self) -> None:
+        """Test clearing callback events."""
+        thread_id = "test-callback-delete"
+
+        # Send a message first
+        SampleTestHelper.post_json(
+            f"{self.base_url}/run",
+            {"message": "Test", "thread_id": thread_id, "wait_for_response": False},
+        )
+
+        # Delete callbacks
+        response = requests.delete(f"{self.base_url}/callbacks/{thread_id}", timeout=TIMEOUT)
+        assert response.status_code == 204
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for Orchestration Chaining Sample
+
+Tests the orchestration chaining sample for sequential agent execution.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite running for durable orchestrations (or Azure Storage account configured)
+
+Usage:
+    # Start Azurite (if not already running)
+    azurite &
+
+    # Run tests
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_04_single_agent_orchestration_chaining.py -v
+"""
+
+import pytest
+
+from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.sample("04_single_agent_orchestration_chaining"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+@pytest.mark.orchestration
+class TestSampleOrchestrationChaining:
+    """Tests for 04_single_agent_orchestration_chaining sample."""
+
+    def test_orchestration_chaining(self, base_url: str) -> None:
+        """Test sequential agent calls in orchestration."""
+        # Start orchestration
+        response = SampleTestHelper.post_json(f"{base_url}/api/singleagent/run", {})
+        assert response.status_code == 202
+        data = response.json()
+        assert "instanceId" in data
+        assert "statusQueryGetUri" in data
+
+        # Wait for completion with output available
+        status = SampleTestHelper.wait_for_orchestration_with_output(data["statusQueryGetUri"])
+        assert status["runtimeStatus"] == "Completed"
+        assert "output" in status
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for MultiAgent Concurrency Sample
+
+Tests the multi-agent concurrency sample for parallel agent execution.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite running for durable orchestrations (or Azure Storage account configured)
+
+Usage:
+    # Start Azurite (if not already running)
+    azurite &
+
+    # Run tests
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_05_multi_agent_orchestration_concurrency.py -v
+"""
+
+import pytest
+
+from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.orchestration,
+    pytest.mark.sample("05_multi_agent_orchestration_concurrency"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+class TestSampleMultiAgentConcurrency:
+    """Tests for 05_multi_agent_orchestration_concurrency sample."""
+
+    def test_concurrent_agents(self, base_url: str) -> None:
+        """Test multiple agents running concurrently."""
+        # Start orchestration
+        response = SampleTestHelper.post_text(f"{base_url}/api/multiagent/run", "What is temperature?")
+        assert response.status_code == 202
+        data = response.json()
+        assert "instanceId" in data
+        assert "statusQueryGetUri" in data
+
+        # Wait for completion
+        status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
+        assert status["runtimeStatus"] == "Completed"
+        output = status["output"]
+        assert "physicist" in output
+        assert "chemist" in output
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for MultiAgent Conditionals Sample
+
+Tests the multi-agent conditionals sample for conditional orchestration logic.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite running for durable orchestrations (or Azure Storage account configured)
+
+Usage:
+    # Start Azurite (if not already running)
+    azurite &
+
+    # Run tests
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_06_multi_agent_orchestration_conditionals.py -v
+"""
+
+import pytest
+
+from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.orchestration,
+    pytest.mark.sample("06_multi_agent_orchestration_conditionals"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+class TestSampleMultiAgentConditionals:
+    """Tests for 06_multi_agent_orchestration_conditionals sample."""
+
+    def test_legitimate_email(self, base_url: str) -> None:
+        """Test conditional logic with legitimate email."""
+        response = SampleTestHelper.post_json(
+            f"{base_url}/api/spamdetection/run",
+            {
+                "email_id": "email-test-001",
+                "email_content": "Hi John, I hope you are doing well. Can you send me the report?",
+            },
+        )
+        assert response.status_code == 202
+        data = response.json()
+        assert "instanceId" in data
+        assert "statusQueryGetUri" in data
+
+        # Wait for completion
+        status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
+        assert status["runtimeStatus"] == "Completed"
+        assert "Email sent:" in status["output"]
+
+    def test_spam_email(self, base_url: str) -> None:
+        """Test conditional logic with spam email."""
+        response = SampleTestHelper.post_json(
+            f"{base_url}/api/spamdetection/run",
+            {"email_id": "email-test-002", "email_content": "URGENT! You have won $1,000,000! Click here now!"},
+        )
+        assert response.status_code == 202
+        data = response.json()
+        assert "instanceId" in data
+
+        # Wait for completion
+        status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
+        assert status["runtimeStatus"] == "Completed"
+        assert "Email marked as spam:" in status["output"]
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,185 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Integration Tests for Human-in-the-Loop (HITL) Orchestration Sample
+
+Tests the HITL orchestration sample for content generation with human approval workflow.
+
+The function app is automatically started by the test fixture.
+
+Prerequisites:
+- Azure OpenAI credentials configured (see packages/azurefunctions/tests/integration_tests/.env.example)
+- Azurite running for durable orchestrations (or Azure Storage account configured)
+
+Usage:
+    # Start Azurite (if not already running)
+    azurite &
+
+    # Run tests
+    uv run pytest packages/azurefunctions/tests/integration_tests/test_07_single_agent_orchestration_hitl.py -v
+"""
+
+import time
+
+import pytest
+
+from .testutils import SampleTestHelper, skip_if_azure_functions_integration_tests_disabled
+
+# Module-level markers - applied to all tests in this file
+pytestmark = [
+    pytest.mark.sample("07_single_agent_orchestration_hitl"),
+    pytest.mark.usefixtures("function_app_for_test"),
+    skip_if_azure_functions_integration_tests_disabled,
+]
+
+
+@pytest.mark.orchestration
+class TestSampleHITLOrchestration:
+    """Tests for 07_single_agent_orchestration_hitl sample."""
+
+    @pytest.fixture(autouse=True)
+    def _set_hitl_base_url(self, base_url: str) -> None:
+        """Prepare the HITL API base URL for the module's tests."""
+        self.hitl_base_url = f"{base_url}/api/hitl"
+
+    def test_hitl_orchestration_approval(self) -> None:
+        """Test HITL orchestration with human approval."""
+        # Start orchestration
+        response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/run",
+            {"topic": "artificial intelligence", "max_review_attempts": 3, "approval_timeout_hours": 1.0},
+        )
+        assert response.status_code == 202
+        data = response.json()
+        assert "instanceId" in data
+        assert "statusQueryGetUri" in data
+        assert data["topic"] == "artificial intelligence"
+        instance_id = data["instanceId"]
+
+        # Wait a bit for the orchestration to generate initial content
+        time.sleep(5)
+
+        # Check status to ensure it's waiting for approval
+        status_response = SampleTestHelper.get(data["statusQueryGetUri"])
+        assert status_response.status_code == 200
+        status = status_response.json()
+        assert status["runtimeStatus"] in ["Running", "Pending"]
+
+        # Send approval
+        approval_response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""}
+        )
+        assert approval_response.status_code == 200
+        approval_data = approval_response.json()
+        assert approval_data["approved"] is True
+
+        # Wait for orchestration to complete
+        status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
+        assert status["runtimeStatus"] == "Completed"
+        assert "output" in status
+        assert "content" in status["output"]
+
+    def test_hitl_orchestration_rejection_with_feedback(self) -> None:
+        """Test HITL orchestration with rejection and subsequent approval."""
+        # Start orchestration
+        response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/run",
+            {"topic": "machine learning", "max_review_attempts": 3, "approval_timeout_hours": 1.0},
+        )
+        assert response.status_code == 202
+        data = response.json()
+        instance_id = data["instanceId"]
+
+        # Wait for initial content generation
+        time.sleep(5)
+
+        # Send rejection with feedback
+        rejection_response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/approve/{instance_id}",
+            {"approved": False, "feedback": "Please make it more concise and focus on practical applications."},
+        )
+        assert rejection_response.status_code == 200
+
+        # Wait for regeneration
+        time.sleep(5)
+
+        # Check status - should still be running
+        status_response = SampleTestHelper.get(data["statusQueryGetUri"])
+        assert status_response.status_code == 200
+        status = status_response.json()
+        assert status["runtimeStatus"] in ["Running", "Pending"]
+
+        # Now approve the revised content
+        approval_response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""}
+        )
+        assert approval_response.status_code == 200
+
+        # Wait for completion
+        status = SampleTestHelper.wait_for_orchestration(data["statusQueryGetUri"])
+        assert status["runtimeStatus"] == "Completed"
+        assert "output" in status
+
+    def test_hitl_orchestration_missing_topic(self) -> None:
+        """Test HITL orchestration with missing topic."""
+        response = SampleTestHelper.post_json(f"{self.hitl_base_url}/run", {"max_review_attempts": 3})
+        assert response.status_code == 400
+        data = response.json()
+        assert "error" in data
+
+    def test_hitl_get_status(self) -> None:
+        """Test getting orchestration status."""
+        # Start orchestration
+        response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/run",
+            {"topic": "quantum computing", "max_review_attempts": 2, "approval_timeout_hours": 1.0},
+        )
+        assert response.status_code == 202
+        data = response.json()
+        instance_id = data["instanceId"]
+
+        # Get status
+        status_response = SampleTestHelper.get(f"{self.hitl_base_url}/status/{instance_id}")
+        assert status_response.status_code == 200
+        status = status_response.json()
+        assert "instanceId" in status
+        assert "runtimeStatus" in status
+        assert status["instanceId"] == instance_id
+
+        # Cleanup: approve to complete orchestration
+        time.sleep(5)
+        SampleTestHelper.post_json(f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""})
+
+    def test_hitl_approval_invalid_payload(self) -> None:
+        """Test sending approval with invalid payload."""
+        # Start orchestration first
+        response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/run",
+            {"topic": "test topic", "max_review_attempts": 1, "approval_timeout_hours": 1.0},
+        )
+        assert response.status_code == 202
+        data = response.json()
+        instance_id = data["instanceId"]
+
+        time.sleep(3)
+
+        # Send approval without 'approved' field
+        approval_response = SampleTestHelper.post_json(
+            f"{self.hitl_base_url}/approve/{instance_id}", {"feedback": "Some feedback"}
+        )
+        assert approval_response.status_code == 400
+        error_data = approval_response.json()
+        assert "error" in error_data
+
+        # Cleanup
+        SampleTestHelper.post_json(f"{self.hitl_base_url}/approve/{instance_id}", {"approved": True, "feedback": ""})
+
+    def test_hitl_status_invalid_instance(self) -> None:
+        """Test getting status for non-existent instance."""
+        response = SampleTestHelper.get(f"{self.hitl_base_url}/status/invalid-instance-id")
+        assert response.status_code == 404
+        data = response.json()
+        assert "error" in data
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -0,0 +1,397 @@
+# Copyright (c) Microsoft. All rights reserved.
+"""
+Shared test helper utilities for sample integration tests.
+
+This module provides common utilities for testing Azure Functions samples.
+"""
+
+import os
+import socket
+import subprocess
+import sys
+import time
+import uuid
+from contextlib import suppress
+from pathlib import Path
+from typing import Any
+
+import pytest
+import requests
+
+# Configuration
+TIMEOUT = 30  # seconds
+ORCHESTRATION_TIMEOUT = 180  # seconds for orchestrations
+_DEFAULT_HOST = "localhost"
+
+
+class FunctionAppStartupError(RuntimeError):
+    """Raised when the Azure Functions host fails to start reliably."""
+
+    pass
+
+
+def _load_env_file_if_present() -> None:
+    """Load environment variables from the local .env file when available."""
+    env_file = Path(__file__).parent / ".env"
+    if not env_file.exists():
+        return
+
+    try:
+        from dotenv import load_dotenv
+
+        load_dotenv(env_file)
+    except ImportError:
+        # python-dotenv not available; rely on existing environment
+        pass
+
+
+def _should_skip_azure_functions_integration_tests() -> tuple[bool, str]:
+    """Determine whether Azure Functions integration tests should be skipped."""
+    _load_env_file_if_present()
+
+    run_integration_tests = os.getenv("RUN_INTEGRATION_TESTS", "false").lower() == "true"
+    if not run_integration_tests:
+        return (
+            True,
+            "Integration tests are disabled. Set RUN_INTEGRATION_TESTS=true to enable Azure Functions sample tests.",
+        )
+
+    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."
+
+    return False, "Integration tests enabled."
+
+
+_SKIP_AZURE_FUNCTIONS_INTEGRATION_TESTS, _AZURE_FUNCTIONS_SKIP_REASON = _should_skip_azure_functions_integration_tests()
+
+skip_if_azure_functions_integration_tests_disabled = pytest.mark.skipif(
+    _SKIP_AZURE_FUNCTIONS_INTEGRATION_TESTS,
+    reason=_AZURE_FUNCTIONS_SKIP_REASON,
+)
+
+
+class SampleTestHelper:
+    """Helper class for testing samples."""
+
+    @staticmethod
+    def post_json(url: str, data: dict[str, Any], timeout: int = TIMEOUT) -> requests.Response:
+        """POST JSON data to a URL."""
+        return requests.post(url, json=data, headers={"Content-Type": "application/json"}, timeout=timeout)
+
+    @staticmethod
+    def post_text(url: str, text: str, timeout: int = TIMEOUT) -> requests.Response:
+        """POST plain text to a URL."""
+        return requests.post(url, data=text, headers={"Content-Type": "text/plain"}, timeout=timeout)
+
+    @staticmethod
+    def get(url: str, timeout: int = TIMEOUT) -> requests.Response:
+        """GET request to a URL."""
+        return requests.get(url, timeout=timeout)
+
+    @staticmethod
+    def wait_for_orchestration(
+        status_url: str, max_wait: int = ORCHESTRATION_TIMEOUT, poll_interval: int = 2
+    ) -> dict[str, Any]:
+        """
+        Wait for an orchestration to complete.
+
+        Args:
+            status_url: URL to poll for orchestration status
+            max_wait: Maximum seconds to wait
+            poll_interval: Seconds between polls
+
+        Returns:
+            Final orchestration status
+
+        Raises:
+            TimeoutError: If orchestration doesn't complete in time
+        """
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            response = requests.get(status_url, timeout=TIMEOUT)
+            response.raise_for_status()
+            status = response.json()
+
+            runtime_status = status.get("runtimeStatus", "")
+            if runtime_status in ["Completed", "Failed", "Terminated"]:
+                return status
+
+            time.sleep(poll_interval)
+
+        raise TimeoutError(f"Orchestration did not complete within {max_wait} seconds")
+
+    @staticmethod
+    def wait_for_orchestration_with_output(
+        status_url: str, max_wait: int = ORCHESTRATION_TIMEOUT, poll_interval: int = 2
+    ) -> dict[str, Any]:
+        """
+        Wait for an orchestration to complete and have output available.
+
+        This is a specialized version of wait_for_orchestration that also
+        ensures the output field is present, handling timing race conditions.
+
+        Args:
+            status_url: URL to poll for orchestration status
+            max_wait: Maximum seconds to wait
+            poll_interval: Seconds between polls
+
+        Returns:
+            Final orchestration status with output
+
+        Raises:
+            TimeoutError: If orchestration doesn't complete with output in time
+        """
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            response = requests.get(status_url, timeout=TIMEOUT)
+            response.raise_for_status()
+            status = response.json()
+
+            runtime_status = status.get("runtimeStatus", "")
+            if runtime_status in ["Failed", "Terminated"]:
+                return status
+            if runtime_status == "Completed" and status.get("output"):
+                return status
+            # If completed but no output, continue polling for a bit more to
+            # handle the race condition where output has not been persisted yet.
+
+            time.sleep(poll_interval)
+
+        # Provide detailed error message based on final status
+        final_response = requests.get(status_url, timeout=TIMEOUT)
+        final_response.raise_for_status()
+        final_status = final_response.json()
+        final_runtime_status = final_status.get("runtimeStatus", "Unknown")
+
+        if final_runtime_status == "Completed":
+            if "output" not in final_status:
+                raise TimeoutError(
+                    "Orchestration completed but 'output' field is missing after "
+                    f"{max_wait} seconds. Final status: {final_status}"
+                )
+            if not final_status["output"]:
+                raise TimeoutError(
+                    "Orchestration completed but output is empty after "
+                    f"{max_wait} seconds. Final status: {final_status}"
+                )
+            raise TimeoutError(
+                "Orchestration completed with output but validation failed after "
+                f"{max_wait} seconds. Final status: {final_status}"
+            )
+        raise TimeoutError(
+            "Orchestration did not complete within "
+            f"{max_wait} seconds. Final status: {final_runtime_status}, "
+            f"Full status: {final_status}"
+        )
+
+
+# Function App Lifecycle Management Helpers
+
+
+def _resolve_repo_root() -> Path:
+    """Resolve the repository root, preferring GITHUB_WORKSPACE when available."""
+    workspace = os.getenv("GITHUB_WORKSPACE")
+    if workspace:
+        candidate = Path(workspace).expanduser()
+        if not (candidate / "samples").exists() and (candidate / "python" / "samples").exists():
+            return (candidate / "python").resolve()
+        return candidate.resolve()
+
+    # If `GITHUB_WORKSPACE` is not set,
+    # go up from testutils.py -> integration_tests -> tests -> azurefunctions -> packages -> python
+    return Path(__file__).resolve().parents[4]
+
+
+def get_sample_path_from_marker(request) -> tuple[Path | None, str | None]:
+    """
+    Get sample path from @pytest.mark.sample() marker.
+
+    Returns a tuple of (sample_path, error_message).
+    If successful, error_message is None.
+    If failed, sample_path is None and error_message contains the reason.
+    """
+    marker = request.node.get_closest_marker("sample")
+
+    if not marker:
+        return (
+            None,
+            (
+                "No @pytest.mark.sample() marker found on test. Add pytestmark with "
+                "@pytest.mark.sample('sample_name') to the test module."
+            ),
+        )
+
+    if not marker.args:
+        return (
+            None,
+            "@pytest.mark.sample() marker found but no sample name provided. Use @pytest.mark.sample('sample_name').",
+        )
+
+    sample_name = marker.args[0]
+    repo_root = _resolve_repo_root()
+    sample_path = repo_root / "samples" / "getting_started" / "azure_functions" / sample_name
+
+    if not sample_path.exists():
+        return None, f"Sample directory does not exist: {sample_path}"
+
+    return sample_path, None
+
+
+def find_available_port(host: str = _DEFAULT_HOST) -> int:
+    """Find an available TCP port on the given host."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind((host, 0))
+        return sock.getsockname()[1]
+
+
+def build_base_url(port: int, host: str = _DEFAULT_HOST) -> str:
+    """Construct a base URL for the Azure Functions host."""
+    return f"http://{host}:{port}"
+
+
+def is_port_in_use(port: int, host: str = _DEFAULT_HOST) -> bool:
+    """
+    Check if a port is already in use.
+
+    Returns True if the port is in use, False otherwise.
+    """
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        return sock.connect_ex((host, port)) == 0
+
+
+def load_and_validate_env() -> None:
+    """
+    Load .env file from current directory if it exists,
+    then validate that required environment variables are present.
+
+    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",
+    ]
+
+    # Check if required env vars are set
+    missing_vars = [var for var in required_env_vars if not os.environ.get(var)]
+
+    if missing_vars:
+        pytest.fail(
+            f"Missing required environment variables: {', '.join(missing_vars)}. "
+            "Please create a .env file in tests/integration_tests/ based on .env.example or "
+            "set these variables in your environment."
+        )
+
+
+def start_function_app(sample_path: Path, port: int) -> subprocess.Popen:
+    """
+    Start a function app in the specified sample directory.
+
+    Returns the subprocess.Popen object for the running process.
+    """
+    env = os.environ.copy()
+    # Use a unique TASKHUB_NAME for each test run to ensure test isolation.
+    # This prevents conflicts between parallel or repeated test runs, as Durable Functions
+    # use the task hub name to separate orchestration state.
+    env["TASKHUB_NAME"] = f"test{uuid.uuid4().hex[:8]}"
+
+    # On Windows, use CREATE_NEW_PROCESS_GROUP to allow proper termination
+    # shell=True only on Windows to handle PATH resolution
+    if sys.platform == "win32":
+        return subprocess.Popen(
+            ["func", "start", "--port", str(port)],
+            cwd=str(sample_path),
+            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
+            shell=True,
+            env=env,
+        )
+    # On Unix, don't use shell=True to avoid shell wrapper issues
+    return subprocess.Popen(["func", "start", "--port", str(port)], cwd=str(sample_path), env=env)
+
+
+def wait_for_function_app_ready(func_process: subprocess.Popen, port: int, max_wait: int = 60) -> None:
+    """Block until the Azure Functions host responds healthy or fail fast."""
+    start_time = time.time()
+    health_url = f"{build_base_url(port)}/api/health"
+    last_error: Exception | None = None
+
+    while time.time() - start_time < max_wait:
+        # If the process exited early, capture any previously seen error and fail fast.
+        if func_process.poll() is not None:
+            raise FunctionAppStartupError(
+                f"Function app process exited with code {func_process.returncode} before becoming healthy"
+            ) from last_error
+
+        if is_port_in_use(port):
+            try:
+                response = requests.get(health_url, timeout=5)
+                if response.status_code == 200:
+                    return
+                last_error = RuntimeError(f"Health check returned {response.status_code}")
+            except requests.RequestException as exc:
+                last_error = exc
+
+        time.sleep(1)
+
+    raise FunctionAppStartupError(
+        f"Function app did not become healthy on port {port} within {max_wait} seconds"
+    ) from last_error
+
+
+def cleanup_function_app(func_process: subprocess.Popen) -> None:
+    """
+    Clean up the function app process and all its children.
+
+    Uses psutil if available for more thorough cleanup, falls back to basic termination.
+    """
+    try:
+        import psutil
+
+        if func_process.poll() is None:  # Process still running
+            # Get parent process
+            parent = psutil.Process(func_process.pid)
+
+            # Get all child processes recursively
+            children = parent.children(recursive=True)
+
+            # Kill children first
+            for child in children:
+                with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
+                    child.kill()
+
+            # Kill parent
+            with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
+                parent.kill()
+
+            # Wait for all to terminate
+            _gone, alive = psutil.wait_procs(children + [parent], timeout=3)
+
+            # Force kill any remaining
+            for proc in alive:
+                with suppress(psutil.NoSuchProcess, psutil.AccessDenied):
+                    proc.kill()
+    except ImportError:
+        # Fallback if psutil not available
+        try:
+            if func_process.poll() is None:
+                func_process.kill()
+                func_process.wait()
+        except Exception:
+            # Ignore all exceptions during fallback cleanup; best effort to terminate process.
+            pass
+    except Exception:
+        pass  # Best effort cleanup
+
+    # Give the port time to be released
+    time.sleep(2)
@@ -0,0 +1,738 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for AgentFunctionApp."""
+
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+from unittest.mock import ANY, AsyncMock, Mock, patch
+
+import azure.durable_functions as df
+import azure.functions as func
+import pytest
+from agent_framework import AgentRunResponse, ChatMessage
+
+from agent_framework_azurefunctions import AgentFunctionApp
+from agent_framework_azurefunctions._app import WAIT_FOR_RESPONSE_FIELD, WAIT_FOR_RESPONSE_HEADER
+from agent_framework_azurefunctions._entities import AgentEntity, AgentState, create_agent_entity
+
+TFunc = TypeVar("TFunc", bound=Callable[..., Any])
+
+
+class TestAgentFunctionAppInit:
+    """Test suite for AgentFunctionApp initialization."""
+
+    def test_init_with_defaults(self) -> None:
+        """Test initialization with default parameters."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent])
+
+        assert len(app.agents) == 1
+        assert "TestAgent" in app.agents
+        assert app.enable_health_check is True
+
+    def test_init_with_custom_auth_level(self) -> None:
+        """Test initialization with custom auth level."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent], http_auth_level=func.AuthLevel.FUNCTION)
+
+        # App should be created successfully
+        assert "TestAgent" in app.agents
+
+    def test_init_with_health_check_disabled(self) -> None:
+        """Test initialization with health check disabled."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent], enable_health_check=False)
+
+        assert app.enable_health_check is False
+
+    def test_init_with_http_endpoints_disabled(self) -> None:
+        """Test initialization with HTTP endpoints disabled."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent], enable_http_endpoints=False)
+
+        assert app.enable_http_endpoints is False
+
+    def test_init_stores_agent_reference(self) -> None:
+        """Test that agent reference is stored correctly."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent])
+
+        assert app.agents["TestAgent"].name == "TestAgent"
+
+    def test_add_agent_uses_specific_callback(self) -> None:
+        """Verify that a per-agent callback overrides the default."""
+
+        mock_agent = Mock()
+        mock_agent.name = "CallbackAgent"
+        specific_callback = Mock()
+
+        with patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_mock:
+            app = AgentFunctionApp(default_callback=Mock())
+            app.add_agent(mock_agent, callback=specific_callback)
+
+        setup_mock.assert_called_once()
+        _, _, passed_callback, enable_http_endpoint = setup_mock.call_args[0]
+        assert passed_callback is specific_callback
+        assert enable_http_endpoint is True
+
+    def test_default_callback_applied_when_no_specific(self) -> None:
+        """Ensure the default callback is supplied when add_agent lacks override."""
+
+        mock_agent = Mock()
+        mock_agent.name = "DefaultAgent"
+        default_callback = Mock()
+
+        with patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_mock:
+            app = AgentFunctionApp(default_callback=default_callback)
+            app.add_agent(mock_agent)
+
+        setup_mock.assert_called_once()
+        _, _, passed_callback, enable_http_endpoint = setup_mock.call_args[0]
+        assert passed_callback is default_callback
+        assert enable_http_endpoint is True
+
+    def test_init_with_agents_uses_default_callback(self) -> None:
+        """Agents provided in __init__ should receive the default callback."""
+
+        mock_agent = Mock()
+        mock_agent.name = "InitAgent"
+        default_callback = Mock()
+
+        with patch.object(AgentFunctionApp, "_setup_agent_functions") as setup_mock:
+            AgentFunctionApp(agents=[mock_agent], default_callback=default_callback)
+
+        setup_mock.assert_called_once()
+        _, _, passed_callback, enable_http_endpoint = setup_mock.call_args[0]
+        assert passed_callback is default_callback
+        assert enable_http_endpoint is True
+
+
+class TestAgentFunctionAppSetup:
+    """Test suite for AgentFunctionApp setup and configuration."""
+
+    def test_app_is_dfapp_instance(self) -> None:
+        """Test that AgentFunctionApp is a DFApp instance."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        app = AgentFunctionApp(agents=[mock_agent])
+
+        assert isinstance(app, df.DFApp)
+
+    def test_setup_creates_http_trigger(self) -> None:
+        """Test that setup creates an HTTP trigger."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        def passthrough_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                return func
+
+            return decorator
+
+        with (
+            patch.object(AgentFunctionApp, "route", new=passthrough_decorator),
+            patch.object(AgentFunctionApp, "durable_client_input", new=passthrough_decorator),
+            patch.object(AgentFunctionApp, "entity_trigger", new=passthrough_decorator),
+        ):
+            app = AgentFunctionApp(agents=[mock_agent])
+
+        # Verify agent is registered
+        assert "TestAgent" in app.agents
+
+    def test_http_function_name_uses_prefix_format(self) -> None:
+        """Ensure function names follow the prefix-agent naming convention."""
+        mock_agent = Mock()
+        mock_agent.name = "Agent 42"
+
+        captured_names: list[str] = []
+
+        def capture_function_name(
+            self: AgentFunctionApp, name: str, *args: Any, **kwargs: Any
+        ) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                captured_names.append(name)
+                return func
+
+            return decorator
+
+        def passthrough_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                return func
+
+            return decorator
+
+        with (
+            patch.object(AgentFunctionApp, "function_name", new=capture_function_name),
+            patch.object(AgentFunctionApp, "route", new=passthrough_decorator),
+            patch.object(AgentFunctionApp, "durable_client_input", new=passthrough_decorator),
+            patch.object(AgentFunctionApp, "entity_trigger", new=passthrough_decorator),
+        ):
+            AgentFunctionApp(agents=[mock_agent])
+
+        assert captured_names == ["http-Agent_42"]
+
+    def test_setup_skips_http_trigger_when_disabled(self) -> None:
+        """Test that HTTP trigger is not created when disabled."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        captured_routes: list[str | None] = []
+
+        def capture_route(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                route_key = kwargs.get("route") if kwargs else None
+                captured_routes.append(route_key)
+                return func
+
+            return decorator
+
+        def passthrough_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                return func
+
+            return decorator
+
+        with (
+            patch.object(AgentFunctionApp, "function_name", new=passthrough_decorator),
+            patch.object(AgentFunctionApp, "route", new=capture_route),
+            patch.object(AgentFunctionApp, "durable_client_input", new=passthrough_decorator),
+            patch.object(AgentFunctionApp, "entity_trigger", new=passthrough_decorator),
+        ):
+            app = AgentFunctionApp(agents=[mock_agent], enable_http_endpoints=False)
+
+        # Verify agent is registered
+        assert "TestAgent" in app.agents
+
+        # Verify that no HTTP run route was created
+        run_route = f"agents/{mock_agent.name}/run"
+        assert run_route not in captured_routes
+
+    def test_agent_override_enables_http_route_when_app_disabled(self) -> None:
+        """Agent-level override should enable HTTP route even when app disables it."""
+
+        mock_agent = Mock()
+        mock_agent.name = "OverrideAgent"
+
+        with (
+            patch.object(AgentFunctionApp, "_setup_http_run_route") as http_route_mock,
+            patch.object(AgentFunctionApp, "_setup_agent_entity") as agent_entity_mock,
+        ):
+            app = AgentFunctionApp(enable_health_check=False, enable_http_endpoints=False)
+            app.add_agent(mock_agent, enable_http_endpoint=True)
+
+        http_route_mock.assert_called_once_with("OverrideAgent")
+        agent_entity_mock.assert_called_once_with(mock_agent, "OverrideAgent", ANY)
+        assert app.agent_http_endpoint_flags["OverrideAgent"] is True
+
+    def test_agent_override_disables_http_route_when_app_enabled(self) -> None:
+        """Agent-level override should disable HTTP route even when app enables it."""
+
+        mock_agent = Mock()
+        mock_agent.name = "DisabledOverride"
+
+        with (
+            patch.object(AgentFunctionApp, "_setup_http_run_route") as http_route_mock,
+            patch.object(AgentFunctionApp, "_setup_agent_entity") as agent_entity_mock,
+        ):
+            app = AgentFunctionApp(enable_health_check=False, enable_http_endpoints=True)
+            app.add_agent(mock_agent, enable_http_endpoint=False)
+
+        http_route_mock.assert_not_called()
+        agent_entity_mock.assert_called_once_with(mock_agent, "DisabledOverride", ANY)
+        assert app.agent_http_endpoint_flags["DisabledOverride"] is False
+
+    def test_multiple_apps_independent(self) -> None:
+        """Test that multiple AgentFunctionApp instances are independent."""
+        agent1 = Mock()
+        agent1.name = "Agent1"
+        agent2 = Mock()
+        agent2.name = "Agent2"
+
+        app1 = AgentFunctionApp(agents=[agent1])
+        app2 = AgentFunctionApp(agents=[agent2])
+
+        assert app1.agents["Agent1"].name == "Agent1"
+        assert app2.agents["Agent2"].name == "Agent2"
+        assert "Agent1" in app1.agents
+        assert "Agent2" in app2.agents
+
+
+class TestWaitForResponseAndCorrelationId:
+    """Tests for wait_for_response flag and correlation ID handling."""
+
+    def _create_app(self) -> AgentFunctionApp:
+        mock_agent = Mock()
+        mock_agent.__class__.__name__ = "MockAgent"
+        mock_agent.name = "MockAgent"
+        return AgentFunctionApp(agents=[mock_agent], enable_health_check=False)
+
+    def _make_request(
+        self,
+        headers: dict[str, str] | None = None,
+        params: dict[str, str] | None = None,
+    ) -> Mock:
+        request = Mock()
+        request.headers = headers or {}
+        request.params = params or {}
+        return request
+
+    def test_wait_for_response_header_true(self) -> None:
+        """Test that the wait-for-response header is honored."""
+        app = self._create_app()
+        request = self._make_request(headers={WAIT_FOR_RESPONSE_HEADER: "true"})
+
+        assert app._should_wait_for_response(request, {}) is True
+
+    def test_wait_for_response_body_snake_case(self) -> None:
+        """Test that payload controls wait_for_response."""
+        app = self._create_app()
+        request = self._make_request()
+
+        assert app._should_wait_for_response(request, {WAIT_FOR_RESPONSE_FIELD: "true"}) is True
+        assert app._should_wait_for_response(request, {WAIT_FOR_RESPONSE_FIELD: "false"}) is False
+        assert app._should_wait_for_response(request, {WAIT_FOR_RESPONSE_FIELD: "0"}) is False
+
+    def test_wait_for_response_query_parameter(self) -> None:
+        """Test that query parameter controls wait_for_response."""
+        app = self._create_app()
+        request = self._make_request(params={WAIT_FOR_RESPONSE_FIELD: "true"})
+
+        assert app._should_wait_for_response(request, {}) is True
+
+    def test_wait_for_response_query_precedence(self) -> None:
+        """Test that query parameter overrides body value."""
+        app = self._create_app()
+        request = self._make_request(params={WAIT_FOR_RESPONSE_FIELD: "false"})
+
+        assert app._should_wait_for_response(request, {WAIT_FOR_RESPONSE_FIELD: "true"}) is False
+
+
+class TestAgentEntityOperations:
+    """Test suite for entity operations."""
+
+    async def test_entity_run_agent_operation(self) -> None:
+        """Test that entity can run agent operation."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(
+            return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Test response")])
+        )
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context,
+            {"message": "Test message", "thread_id": "test-conv-123", "correlation_id": "corr-app-entity-1"},
+        )
+
+        assert result["status"] == "success"
+        assert result["response"] == "Test response"
+        assert result["message"] == "Test message"
+        assert result["thread_id"] == "test-conv-123"
+        assert entity.state.message_count == 1
+
+    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=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response 1")])
+        )
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send first message
+        await entity.run_agent(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-app-entity-2"}
+        )
+
+        history = entity.state.conversation_history
+        assert len(history) == 2  # User + assistant
+
+        user_msg = history[0]
+        user_role = getattr(user_msg.role, "value", user_msg.role)
+        assert user_role == "user"
+        assert user_msg.text == "Message 1"
+
+        assistant_msg = history[1]
+        assistant_role = getattr(assistant_msg.role, "value", assistant_msg.role)
+        assert assistant_role == "assistant"
+        assert assistant_msg.text == "Response 1"
+
+    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=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response")])
+        )
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        assert entity.state.message_count == 0
+
+        await entity.run_agent(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-app-entity-3a"}
+        )
+        assert entity.state.message_count == 1
+
+        await entity.run_agent(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-app-entity-3b"}
+        )
+        assert entity.state.message_count == 2
+
+    def test_entity_reset(self) -> None:
+        """Test that entity reset clears state."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        # Set some state
+        entity.state.message_count = 10
+        entity.state.last_response = "Some response"
+        entity.state.conversation_history = [
+            ChatMessage(role="user", text="test", additional_properties={"timestamp": "2024-01-01T00:00:00Z"})
+        ]
+
+        # Reset
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert entity.state.message_count == 0
+        assert entity.state.last_response is None
+        assert len(entity.state.conversation_history) == 0
+
+
+class TestAgentEntityFactory:
+    """Test suite for the entity factory function."""
+
+    def test_create_agent_entity_returns_function(self) -> None:
+        """Test that create_agent_entity returns a function."""
+        mock_agent = Mock()
+        entity_function = create_agent_entity(mock_agent)
+
+        assert callable(entity_function)
+
+    def test_entity_function_handles_run_agent_operation(self) -> None:
+        """Test that the entity function handles the run_agent operation."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(
+            return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response")])
+        )
+
+        entity_function = create_agent_entity(mock_agent)
+
+        # Mock context
+        mock_context = Mock()
+        mock_context.operation_name = "run_agent"
+        mock_context.get_input.return_value = {
+            "message": "Test message",
+            "thread_id": "conv-123",
+            "correlation_id": "corr-app-factory-1",
+        }
+        mock_context.get_state.return_value = None
+
+        # Execute entity function
+        entity_function(mock_context)
+
+        # Verify result was set
+        assert mock_context.set_result.called
+        assert mock_context.set_state.called
+
+    def test_entity_function_handles_reset_operation(self) -> None:
+        """Test that the entity function handles the reset operation."""
+        mock_agent = Mock()
+        entity_function = create_agent_entity(mock_agent)
+
+        # Mock context
+        mock_context = Mock()
+        mock_context.operation_name = "reset"
+        mock_context.get_state.return_value = {
+            "message_count": 5,
+            "conversation_history": [{"role": "user", "content": "test"}],
+            "last_response": "Test",
+        }
+
+        # Execute entity function
+        entity_function(mock_context)
+
+        # Verify result was set
+        assert mock_context.set_result.called
+        result_call = mock_context.set_result.call_args[0][0]
+        assert result_call["status"] == "reset"
+
+    def test_entity_function_handles_unknown_operation(self) -> None:
+        """Test that the entity function handles an unknown operation."""
+        mock_agent = Mock()
+        entity_function = create_agent_entity(mock_agent)
+
+        # Mock context with unknown operation
+        mock_context = Mock()
+        mock_context.operation_name = "unknown_operation"
+        mock_context.get_state.return_value = None
+
+        # Execute entity function
+        entity_function(mock_context)
+
+        # Verify error result was set
+        assert mock_context.set_result.called
+        result_call = mock_context.set_result.call_args[0][0]
+        assert "error" in result_call
+        assert "unknown_operation" in result_call["error"]
+
+    def test_entity_function_restores_state(self) -> None:
+        """Test that the entity function restores state from the context."""
+        mock_agent = Mock()
+        entity_function = create_agent_entity(mock_agent)
+
+        # Mock context with existing state
+        existing_state = {
+            "message_count": 3,
+            "conversation_history": [{"role": "user", "content": "msg1"}, {"role": "assistant", "content": "resp1"}],
+            "last_response": "resp1",
+        }
+
+        mock_context = Mock()
+        mock_context.operation_name = "reset"
+        mock_context.get_state.return_value = existing_state
+
+        with patch.object(AgentState, "restore_state") as restore_state_mock:
+            entity_function(mock_context)
+
+        restore_state_mock.assert_called_once_with(existing_state)
+
+
+class TestErrorHandling:
+    """Test suite for error handling."""
+
+    async def test_entity_handles_agent_error(self) -> None:
+        """Test that the entity handles agent execution errors."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=Exception("Agent error"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Test message", "thread_id": "conv-1", "correlation_id": "corr-app-error-1"}
+        )
+
+        assert result["status"] == "error"
+        assert "error" in result
+        assert "Agent error" in result["error"]
+        assert result["error_type"] == "Exception"
+
+    def test_entity_function_handles_exception(self) -> None:
+        """Test that the entity function handles exceptions gracefully."""
+        mock_agent = Mock()
+        # Force an exception by making get_input fail
+        mock_agent.run = AsyncMock(side_effect=Exception("Test error"))
+
+        entity_function = create_agent_entity(mock_agent)
+
+        mock_context = Mock()
+        mock_context.operation_name = "run_agent"
+        mock_context.get_input.side_effect = Exception("Input error")
+        mock_context.get_state.return_value = None
+
+        # Execute entity function - should not raise
+        entity_function(mock_context)
+
+        # Verify error result was set
+        assert mock_context.set_result.called
+        result_call = mock_context.set_result.call_args[0][0]
+        assert "error" in result_call
+
+
+class TestIncomingRequestParsing:
+    """Tests for parsing run requests with JSON and plain text bodies."""
+
+    def _create_app(self) -> AgentFunctionApp:
+        mock_agent = Mock()
+        mock_agent.name = "ParserAgent"
+        return AgentFunctionApp(agents=[mock_agent], enable_health_check=False)
+
+    def test_parse_plain_text_body(self) -> None:
+        """Test parsing a plain-text request body."""
+        app = self._create_app()
+
+        request = Mock()
+        request.headers = {}
+        request.params = {}
+        request.get_json.side_effect = ValueError("Invalid JSON")
+        request.get_body.return_value = b"Plain text message"
+
+        req_body, message, response_format = app._parse_incoming_request(request)
+
+        assert req_body == {}
+        assert message == "Plain text message"
+
+        assert response_format == "text"
+
+    def test_parse_plain_text_trims_whitespace(self) -> None:
+        """Plain-text parser returns an empty string when the body contains only whitespace."""
+        app = self._create_app()
+
+        request = Mock()
+        request.headers = {}
+        request.params = {}
+        request.get_json.side_effect = ValueError("Invalid JSON")
+        request.get_body.return_value = b"   "
+
+        req_body, message, response_format = app._parse_incoming_request(request)
+
+        assert req_body == {}
+        assert message == ""
+        assert response_format == "text"
+
+    def test_accept_header_prefers_json(self) -> None:
+        """Test that the Accept header can force JSON responses for plain-text bodies."""
+        app = self._create_app()
+
+        request = Mock()
+        request.headers = {"accept": "application/json"}
+        request.params = {}
+        request.get_json.side_effect = ValueError("Invalid JSON")
+        request.get_body.return_value = b"Plain text message"
+
+        _, message, response_format = app._parse_incoming_request(request)
+
+        assert message == "Plain text message"
+        assert response_format == "json"
+
+    def test_extract_thread_id_from_query_params(self) -> None:
+        """Test thread identifier extraction from query parameters."""
+        app = self._create_app()
+
+        request = Mock()
+        request.params = {"thread_id": "query-thread"}
+        req_body = {}
+
+        thread_id = app._resolve_thread_id(request, req_body)
+
+        assert thread_id == "query-thread"
+
+
+class TestHttpRunRoute:
+    """Tests for the HTTP run route behavior."""
+
+    @staticmethod
+    def _get_run_handler(agent: Mock) -> Callable[[func.HttpRequest, Any], Awaitable[func.HttpResponse]]:
+        captured_handlers: dict[str | None, Callable[..., Awaitable[func.HttpResponse]]] = {}
+
+        def capture_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                return func
+
+            return decorator
+
+        def capture_route(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
+            def decorator(func: TFunc) -> TFunc:
+                route_key = kwargs.get("route") if kwargs else None
+                captured_handlers[route_key] = func
+                return func
+
+            return decorator
+
+        with (
+            patch.object(AgentFunctionApp, "function_name", new=capture_decorator),
+            patch.object(AgentFunctionApp, "route", new=capture_route),
+            patch.object(AgentFunctionApp, "durable_client_input", new=capture_decorator),
+            patch.object(AgentFunctionApp, "entity_trigger", new=capture_decorator),
+        ):
+            AgentFunctionApp(agents=[agent], enable_health_check=False)
+
+        run_route = f"agents/{agent.name}/run"
+        return captured_handlers[run_route]
+
+    async def test_http_run_accepts_plain_text(self) -> None:
+        """Test that the HTTP handler accepts plain-text requests."""
+        mock_agent = Mock()
+        mock_agent.name = "HttpAgent"
+
+        handler = self._get_run_handler(mock_agent)
+
+        request = Mock()
+        request.headers = {WAIT_FOR_RESPONSE_HEADER: "false"}
+        request.params = {}
+        request.route_params = {}
+        request.get_json.side_effect = ValueError("Invalid JSON")
+        request.get_body.return_value = b"Plain text via HTTP"
+
+        client = AsyncMock()
+
+        response = await handler(request, client)
+
+        assert response.status_code == 202
+        assert response.mimetype == "text/plain"
+        assert response.headers.get("x-ms-thread-id") is not None
+        assert response.get_body().decode("utf-8") == "Agent request accepted"
+
+        signal_args = client.signal_entity.call_args[0]
+        run_request = signal_args[2]
+
+        assert run_request["message"] == "Plain text via HTTP"
+        assert run_request["role"] == "user"
+        assert "thread_id" in run_request
+
+    async def test_http_run_accept_header_returns_json(self) -> None:
+        """Test that Accept header requesting JSON results in JSON response."""
+        mock_agent = Mock()
+        mock_agent.name = "HttpAgentJson"
+
+        handler = self._get_run_handler(mock_agent)
+
+        request = Mock()
+        request.headers = {WAIT_FOR_RESPONSE_HEADER: "false", "Accept": "application/json"}
+        request.params = {}
+        request.route_params = {}
+        request.get_json.side_effect = ValueError("Invalid JSON")
+        request.get_body.return_value = b"Plain text via HTTP"
+
+        client = AsyncMock()
+
+        response = await handler(request, client)
+
+        assert response.status_code == 202
+        assert response.mimetype == "application/json"
+        assert response.headers.get("x-ms-thread-id") is None
+        body = response.get_body().decode("utf-8")
+        assert '"status": "accepted"' in body
+
+    async def test_http_run_rejects_empty_message(self) -> None:
+        """Test that the HTTP handler rejects empty messages with a 400 response."""
+        mock_agent = Mock()
+        mock_agent.name = "HttpAgentEmpty"
+
+        handler = self._get_run_handler(mock_agent)
+
+        request = Mock()
+        request.headers = {WAIT_FOR_RESPONSE_HEADER: "false"}
+        request.params = {}
+        request.route_params = {}
+        request.get_json.side_effect = ValueError("Invalid JSON")
+        request.get_body.return_value = b"   "
+
+        client = AsyncMock()
+
+        response = await handler(request, client)
+
+        assert response.status_code == 400
+        assert response.mimetype == "text/plain"
+        assert response.headers.get("x-ms-thread-id") is not None
+        assert response.get_body().decode("utf-8") == "Message is required"
+        client.signal_entity.assert_not_called()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])
@@ -0,0 +1,904 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for AgentEntity and entity operations.
+
+Run with: pytest tests/test_entities.py -v
+"""
+
+import asyncio
+from collections.abc import AsyncIterator, Callable
+from datetime import datetime
+from typing import Any, TypeVar
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+from agent_framework import AgentRunResponse, AgentRunResponseUpdate, ChatMessage, Role
+from pydantic import BaseModel
+
+from agent_framework_azurefunctions._entities import AgentEntity, create_agent_entity
+from agent_framework_azurefunctions._models import RunRequest
+from agent_framework_azurefunctions._state import AgentState
+
+TFunc = TypeVar("TFunc", bound=Callable[..., Any])
+
+
+def _role_value(chat_message: ChatMessage) -> str:
+    """Helper to extract the string role from a ChatMessage."""
+    role = getattr(chat_message, "role", None)
+    role_value = getattr(role, "value", role)
+    if role_value is None:
+        return ""
+    return str(role_value)
+
+
+def _agent_response(text: str | None) -> AgentRunResponse:
+    """Create an AgentRunResponse with a single assistant message."""
+    message = (
+        ChatMessage(role="assistant", text=text) if text is not None else ChatMessage(role="assistant", contents=[])
+    )
+    return AgentRunResponse(messages=[message])
+
+
+class RecordingCallback:
+    """Callback implementation capturing streaming and final responses for assertions."""
+
+    def __init__(self):
+        self.stream_mock = AsyncMock()
+        self.response_mock = AsyncMock()
+
+    async def on_streaming_response_update(
+        self,
+        update: AgentRunResponseUpdate,
+        context: Any,
+    ) -> None:
+        await self.stream_mock(update, context)
+
+    async def on_agent_response(self, response: AgentRunResponse, context: Any) -> None:
+        await self.response_mock(response, context)
+
+
+class EntityStructuredResponse(BaseModel):
+    answer: float
+
+
+class TestAgentEntityInit:
+    """Test suite for AgentEntity initialization."""
+
+    def test_init_creates_entity(self) -> None:
+        """Test that AgentEntity initializes correctly."""
+        mock_agent = Mock()
+
+        entity = AgentEntity(mock_agent)
+
+        assert entity.agent == mock_agent
+        assert entity.state.conversation_history == []
+        assert entity.state.last_response is None
+        assert entity.state.message_count == 0
+
+    def test_init_stores_agent_reference(self) -> None:
+        """Test that the agent reference is stored correctly."""
+        mock_agent = Mock()
+        mock_agent.name = "TestAgent"
+
+        entity = AgentEntity(mock_agent)
+
+        assert entity.agent.name == "TestAgent"
+
+    def test_init_with_different_agent_types(self) -> None:
+        """Test initialization with different agent types."""
+        agent1 = Mock()
+        agent1.__class__.__name__ = "AzureOpenAIAgent"
+
+        agent2 = Mock()
+        agent2.__class__.__name__ = "CustomAgent"
+
+        entity1 = AgentEntity(agent1)
+        entity2 = AgentEntity(agent2)
+
+        assert entity1.agent.__class__.__name__ == "AzureOpenAIAgent"
+        assert entity2.agent.__class__.__name__ == "CustomAgent"
+
+
+class TestAgentEntityRunAgent:
+    """Test suite for the run_agent operation."""
+
+    async def test_run_agent_executes_agent(self) -> None:
+        """Test that run_agent executes the agent."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Test response")
+        mock_agent.run = AsyncMock(return_value=mock_response)
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Test message", "thread_id": "conv-123", "correlation_id": "corr-entity-1"}
+        )
+
+        # Verify agent.run was called
+        mock_agent.run.assert_called_once()
+        _, kwargs = mock_agent.run.call_args
+        sent_messages = kwargs.get("messages")
+        assert isinstance(sent_messages, list)
+        assert len(sent_messages) == 1
+        sent_message = sent_messages[0]
+        assert isinstance(sent_message, ChatMessage)
+        assert sent_message.text == "Test message"
+        assert _role_value(sent_message) == "user"
+
+        # Verify result
+        assert result["status"] == "success"
+        assert result["response"] == "Test response"
+        assert result["message"] == "Test message"
+        assert result["thread_id"] == "conv-123"
+
+    async def test_run_agent_streaming_callbacks_invoked(self) -> None:
+        """Ensure streaming updates trigger callbacks and run() is not used."""
+
+        updates = [
+            AgentRunResponseUpdate(text="Hello"),
+            AgentRunResponseUpdate(text=" world"),
+        ]
+
+        async def update_generator() -> AsyncIterator[AgentRunResponseUpdate]:
+            for update in updates:
+                yield update
+
+        mock_agent = Mock()
+        mock_agent.name = "StreamingAgent"
+        mock_agent.run_stream = Mock(return_value=update_generator())
+        mock_agent.run = AsyncMock(side_effect=AssertionError("run() should not be called when streaming succeeds"))
+
+        callback = RecordingCallback()
+        entity = AgentEntity(mock_agent, callback=callback)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context,
+            {
+                "message": "Tell me something",
+                "thread_id": "session-1",
+                "correlation_id": "corr-stream-1",
+            },
+        )
+
+        assert result["status"] == "success"
+        assert "Hello" in result.get("response", "")
+        assert callback.stream_mock.await_count == len(updates)
+        assert callback.response_mock.await_count == 1
+        mock_agent.run.assert_not_called()
+
+        # Validate callback arguments
+        stream_calls = callback.stream_mock.await_args_list
+        for expected_update, recorded_call in zip(updates, stream_calls, strict=True):
+            assert recorded_call.args[0] is expected_update
+            context = recorded_call.args[1]
+            assert context.agent_name == "StreamingAgent"
+            assert context.correlation_id == "corr-stream-1"
+            assert context.thread_id == "session-1"
+            assert context.request_message == "Tell me something"
+
+        final_call = callback.response_mock.await_args
+        assert final_call is not None
+        final_response, final_context = final_call.args
+        assert final_context.agent_name == "StreamingAgent"
+        assert final_context.correlation_id == "corr-stream-1"
+        assert final_context.thread_id == "session-1"
+        assert final_context.request_message == "Tell me something"
+        assert getattr(final_response, "text", "").strip()
+
+    async def test_run_agent_final_callback_without_streaming(self) -> None:
+        """Ensure the final callback fires even when streaming is unavailable."""
+
+        mock_agent = Mock()
+        mock_agent.name = "NonStreamingAgent"
+        mock_agent.run_stream = None
+        agent_response = _agent_response("Final response")
+        mock_agent.run = AsyncMock(return_value=agent_response)
+
+        callback = RecordingCallback()
+        entity = AgentEntity(mock_agent, callback=callback)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context,
+            {
+                "message": "Hi",
+                "thread_id": "session-2",
+                "correlation_id": "corr-final-1",
+            },
+        )
+
+        assert result["status"] == "success"
+        assert result.get("response") == "Final response"
+        assert callback.stream_mock.await_count == 0
+        assert callback.response_mock.await_count == 1
+
+        final_call = callback.response_mock.await_args
+        assert final_call is not None
+        assert final_call.args[0] is agent_response
+        final_context = final_call.args[1]
+        assert final_context.agent_name == "NonStreamingAgent"
+        assert final_context.correlation_id == "corr-final-1"
+        assert final_context.thread_id == "session-2"
+        assert final_context.request_message == "Hi"
+
+    async def test_run_agent_updates_conversation_history(self) -> None:
+        """Test that run_agent updates the conversation history."""
+        mock_agent = Mock()
+        mock_response = _agent_response("Agent response")
+        mock_agent.run = AsyncMock(return_value=mock_response)
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run_agent(
+            mock_context, {"message": "User message", "thread_id": "conv-1", "correlation_id": "corr-entity-2"}
+        )
+
+        # Should have 2 entries: user message + assistant response
+        history = entity.state.conversation_history
+
+        assert len(history) == 2
+
+        user_msg = history[0]
+        assert _role_value(user_msg) == "user"
+        assert user_msg.text == "User message"
+
+        assistant_msg = history[1]
+        assert _role_value(assistant_msg) == "assistant"
+        assert assistant_msg.text == "Agent response"
+
+    async def test_run_agent_increments_message_count(self) -> None:
+        """Test that run_agent increments the message count."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        assert entity.state.message_count == 0
+
+        await entity.run_agent(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-entity-3a"}
+        )
+        assert entity.state.message_count == 1
+
+        await entity.run_agent(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-entity-3b"}
+        )
+        assert entity.state.message_count == 2
+
+        await entity.run_agent(
+            mock_context, {"message": "Message 3", "thread_id": "conv-1", "correlation_id": "corr-entity-3c"}
+        )
+        assert entity.state.message_count == 3
+
+    async def test_run_agent_stores_last_response(self) -> None:
+        """Test that run_agent stores the last response."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 1"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run_agent(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-entity-4a"}
+        )
+        assert entity.state.last_response == "Response 1"
+
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 2"))
+        await entity.run_agent(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-entity-4b"}
+        )
+        assert entity.state.last_response == "Response 2"
+
+    async def test_run_agent_with_none_thread_id(self) -> None:
+        """Test run_agent with a None thread identifier."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        with pytest.raises(ValueError, match="thread_id"):
+            await entity.run_agent(
+                mock_context, {"message": "Message", "thread_id": None, "correlation_id": "corr-entity-5"}
+            )
+
+    async def test_run_agent_handles_response_without_text_attribute(self) -> None:
+        """Test that run_agent handles responses without a text attribute."""
+        mock_agent = Mock()
+
+        class NoTextResponse(AgentRunResponse):
+            @property
+            def text(self) -> str:  # type: ignore[override]
+                raise AttributeError("text attribute missing")
+
+        mock_response = NoTextResponse(messages=[ChatMessage(role="assistant", text="ignored")])
+        mock_agent.run = AsyncMock(return_value=mock_response)
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlation_id": "corr-entity-6"}
+        )
+
+        # Should handle gracefully
+        assert result["status"] == "success"
+        assert result["response"] == "Error extracting response"
+
+    async def test_run_agent_handles_none_response_text(self) -> None:
+        """Test that run_agent handles responses with None text."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response(None))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlation_id": "corr-entity-7"}
+        )
+
+        assert result["status"] == "success"
+        assert result["response"] == "No response"
+
+    async def test_run_agent_multiple_conversations(self) -> None:
+        """Test that run_agent maintains history across multiple messages."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send multiple messages
+        await entity.run_agent(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-entity-8a"}
+        )
+        await entity.run_agent(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-entity-8b"}
+        )
+        await entity.run_agent(
+            mock_context, {"message": "Message 3", "thread_id": "conv-1", "correlation_id": "corr-entity-8c"}
+        )
+
+        history = entity.state.conversation_history
+        assert len(history) == 6
+        assert entity.state.message_count == 3
+
+
+class TestAgentEntityReset:
+    """Test suite for the reset operation."""
+
+    def test_reset_clears_conversation_history(self) -> None:
+        """Test that reset clears the conversation history."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        # Add some history
+        entity.state.conversation_history = [
+            ChatMessage(role="user", text="msg1"),
+            ChatMessage(role="assistant", text="resp1"),
+        ]
+
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert entity.state.conversation_history == []
+
+    def test_reset_clears_last_response(self) -> None:
+        """Test that reset clears the last response."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        entity.state.last_response = "Some response"
+
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert entity.state.last_response is None
+
+    def test_reset_clears_message_count(self) -> None:
+        """Test that reset clears the message count."""
+        mock_agent = Mock()
+        entity = AgentEntity(mock_agent)
+
+        entity.state.message_count = 10
+
+        mock_context = Mock()
+        entity.reset(mock_context)
+
+        assert entity.state.message_count == 0
+
+    async def test_reset_after_conversation(self) -> None:
+        """Test reset after a full conversation."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Have a conversation
+        await entity.run_agent(
+            mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-entity-10a"}
+        )
+        await entity.run_agent(
+            mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-entity-10b"}
+        )
+
+        # Verify state before reset
+        assert entity.state.message_count == 2
+        assert len(entity.state.conversation_history) == 4
+
+        # Reset
+        entity.reset(mock_context)
+
+        # Verify state after reset
+        assert entity.state.message_count == 0
+        assert len(entity.state.conversation_history) == 0
+        assert entity.state.last_response is None
+
+
+class TestCreateAgentEntity:
+    """Test suite for the create_agent_entity factory function."""
+
+    def test_create_agent_entity_returns_callable(self) -> None:
+        """Test that create_agent_entity returns a callable."""
+        mock_agent = Mock()
+
+        entity_function = create_agent_entity(mock_agent)
+
+        assert callable(entity_function)
+
+    def test_entity_function_handles_run_agent(self) -> None:
+        """Test that the entity function handles the run_agent operation."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity_function = create_agent_entity(mock_agent)
+
+        # Mock context
+        mock_context = Mock()
+        mock_context.operation_name = "run_agent"
+        mock_context.get_input.return_value = {
+            "message": "Test message",
+            "thread_id": "conv-123",
+            "correlation_id": "corr-entity-factory",
+        }
+        mock_context.get_state.return_value = None
+
+        # Execute
+        entity_function(mock_context)
+
+        # Verify result and state were set
+        assert mock_context.set_result.called
+        assert mock_context.set_state.called
+
+    def test_entity_function_handles_reset(self) -> None:
+        """Test that the entity function handles the reset operation."""
+        mock_agent = Mock()
+
+        entity_function = create_agent_entity(mock_agent)
+
+        # Mock context with existing state
+        mock_context = Mock()
+        mock_context.operation_name = "reset"
+        mock_context.get_state.return_value = {
+            "message_count": 5,
+            "conversation_history": [
+                ChatMessage(
+                    role="user", text="test", additional_properties={"timestamp": "2024-01-01T00:00:00Z"}
+                ).to_dict()
+            ],
+            "last_response": "Test",
+        }
+
+        # Execute
+        entity_function(mock_context)
+
+        # Verify reset result
+        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert result["status"] == "reset"
+
+        # Verify state was cleared
+        assert mock_context.set_state.called
+        state = mock_context.set_state.call_args[0][0]
+        assert state["message_count"] == 0
+        assert state["conversation_history"] == []
+        assert state["last_response"] is None
+
+    def test_entity_function_handles_unknown_operation(self) -> None:
+        """Test that the entity function handles unknown operations."""
+        mock_agent = Mock()
+
+        entity_function = create_agent_entity(mock_agent)
+
+        mock_context = Mock()
+        mock_context.operation_name = "invalid_operation"
+        mock_context.get_state.return_value = None
+
+        # Execute
+        entity_function(mock_context)
+
+        # Verify error result
+        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert "error" in result
+        assert "invalid_operation" in result["error"].lower()
+
+    def test_entity_function_creates_new_entity_on_first_call(self) -> None:
+        """Test that the entity function creates a new entity when no state exists."""
+        mock_agent = Mock()
+        mock_agent.__class__.__name__ = "Agent"
+
+        entity_function = create_agent_entity(mock_agent)
+        mock_context = Mock()
+        mock_context.operation_name = "reset"
+        mock_context.get_state.return_value = None  # No existing state
+
+        # Execute
+        entity_function(mock_context)
+
+        # Verify new entity state was created
+        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert result["status"] == "reset"
+        assert mock_context.set_state.called
+        state = mock_context.set_state.call_args[0][0]
+        assert state["message_count"] == 0
+        assert state["conversation_history"] == []
+
+    def test_entity_function_restores_existing_state(self) -> None:
+        """Test that the entity function restores existing state."""
+        mock_agent = Mock()
+
+        entity_function = create_agent_entity(mock_agent)
+
+        existing_state = {
+            "message_count": 5,
+            "conversation_history": [
+                ChatMessage(
+                    role="user", text="msg1", additional_properties={"timestamp": "2024-01-01T00:00:00Z"}
+                ).to_dict(),
+                ChatMessage(
+                    role="assistant", text="resp1", additional_properties={"timestamp": "2024-01-01T00:05:00Z"}
+                ).to_dict(),
+            ],
+            "last_response": "resp1",
+        }
+
+        mock_context = Mock()
+        mock_context.operation_name = "reset"
+        mock_context.get_state.return_value = existing_state
+
+        with patch.object(AgentState, "restore_state") as restore_state_mock:
+            entity_function(mock_context)
+
+        restore_state_mock.assert_called_once_with(existing_state)
+
+
+class TestErrorHandling:
+    """Test suite for error handling in entities."""
+
+    async def test_run_agent_handles_agent_exception(self) -> None:
+        """Test that run_agent handles agent exceptions."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=Exception("Agent failed"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlation_id": "corr-entity-error-1"}
+        )
+
+        assert result["status"] == "error"
+        assert "error" in result
+        assert "Agent failed" in result["error"]
+        assert result["error_type"] == "Exception"
+
+    async def test_run_agent_handles_value_error(self) -> None:
+        """Test that run_agent handles ValueError instances."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=ValueError("Invalid input"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlation_id": "corr-entity-error-2"}
+        )
+
+        assert result["status"] == "error"
+        assert result["error_type"] == "ValueError"
+        assert "Invalid input" in result["error"]
+
+    async def test_run_agent_handles_timeout_error(self) -> None:
+        """Test that run_agent handles TimeoutError instances."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=TimeoutError("Request timeout"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlation_id": "corr-entity-error-3"}
+        )
+
+        assert result["status"] == "error"
+        assert result["error_type"] == "TimeoutError"
+
+    def test_entity_function_handles_exception_in_operation(self) -> None:
+        """Test that the entity function handles exceptions gracefully."""
+        mock_agent = Mock()
+
+        entity_function = create_agent_entity(mock_agent)
+
+        mock_context = Mock()
+        mock_context.operation_name = "run_agent"
+        mock_context.get_input.side_effect = Exception("Input error")
+        mock_context.get_state.return_value = None
+
+        # Execute - should not raise
+        entity_function(mock_context)
+
+        # Verify error was set
+        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert "error" in result
+
+    async def test_run_agent_preserves_message_on_error(self) -> None:
+        """Test that run_agent preserves message information on error."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(side_effect=Exception("Error"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        result = await entity.run_agent(
+            mock_context,
+            {"message": "Test message", "thread_id": "conv-123", "correlation_id": "corr-entity-error-4"},
+        )
+
+        # Even on error, message info should be preserved
+        assert result["message"] == "Test message"
+        assert result["thread_id"] == "conv-123"
+        assert result["status"] == "error"
+
+
+class TestConversationHistory:
+    """Test suite for conversation history tracking."""
+
+    async def test_conversation_history_has_timestamps(self) -> None:
+        """Test that conversation history entries include timestamps."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run_agent(
+            mock_context, {"message": "Message", "thread_id": "conv-1", "correlation_id": "corr-entity-history-1"}
+        )
+
+        # Check both user and assistant messages have timestamps
+        for entry in entity.state.conversation_history:
+            timestamp = entry.additional_properties.get("timestamp")
+            assert timestamp is not None
+            # Verify timestamp is in ISO format
+            datetime.fromisoformat(timestamp)
+
+    async def test_conversation_history_ordering(self) -> None:
+        """Test that conversation history maintains the correct order."""
+        mock_agent = Mock()
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send multiple messages with different responses
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 1"))
+        await entity.run_agent(
+            mock_context,
+            {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-entity-history-2a"},
+        )
+
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 2"))
+        await entity.run_agent(
+            mock_context,
+            {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-entity-history-2b"},
+        )
+
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response 3"))
+        await entity.run_agent(
+            mock_context,
+            {"message": "Message 3", "thread_id": "conv-1", "correlation_id": "corr-entity-history-2c"},
+        )
+
+        # Verify order
+        history = entity.state.conversation_history
+        assert history[0].text == "Message 1"
+        assert history[1].text == "Response 1"
+        assert history[2].text == "Message 2"
+        assert history[3].text == "Response 2"
+        assert history[4].text == "Message 3"
+        assert history[5].text == "Response 3"
+
+    async def test_conversation_history_role_alternation(self) -> None:
+        """Test that conversation history alternates between user and assistant roles."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        await entity.run_agent(
+            mock_context,
+            {"message": "Message 1", "thread_id": "conv-1", "correlation_id": "corr-entity-history-3a"},
+        )
+        await entity.run_agent(
+            mock_context,
+            {"message": "Message 2", "thread_id": "conv-1", "correlation_id": "corr-entity-history-3b"},
+        )
+
+        # Check role alternation
+        history = entity.state.conversation_history
+        assert _role_value(history[0]) == "user"
+        assert _role_value(history[1]) == "assistant"
+        assert _role_value(history[2]) == "user"
+        assert _role_value(history[3]) == "assistant"
+
+
+class TestRunRequestSupport:
+    """Test suite for RunRequest support in entities."""
+
+    async def test_run_agent_with_run_request_object(self) -> None:
+        """Test run_agent with a RunRequest object."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request = RunRequest(
+            message="Test message",
+            thread_id="conv-123",
+            role=Role.USER,
+            enable_tool_calls=True,
+            correlation_id="corr-runreq-1",
+        )
+
+        result = await entity.run_agent(mock_context, request)
+
+        assert result["status"] == "success"
+        assert result["response"] == "Response"
+        assert result["message"] == "Test message"
+        assert result["thread_id"] == "conv-123"
+
+    async def test_run_agent_with_dict_request(self) -> None:
+        """Test run_agent with a dictionary request."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request_dict = {
+            "message": "Test message",
+            "thread_id": "conv-456",
+            "role": "system",
+            "enable_tool_calls": False,
+            "correlation_id": "corr-runreq-2",
+        }
+
+        result = await entity.run_agent(mock_context, request_dict)
+
+        assert result["status"] == "success"
+        assert result["message"] == "Test message"
+        assert result["thread_id"] == "conv-456"
+
+    async def test_run_agent_with_string_raises_without_correlation(self) -> None:
+        """Test that run_agent rejects legacy string input without correlation ID."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        with pytest.raises(ValueError):
+            await entity.run_agent(mock_context, "Simple message")
+
+    async def test_run_agent_stores_role_in_history(self) -> None:
+        """Test that run_agent stores the role in conversation history."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        # Send as system role
+        request = RunRequest(
+            message="System message",
+            thread_id="conv-runreq-3",
+            role=Role.SYSTEM,
+            correlation_id="corr-runreq-3",
+        )
+
+        await entity.run_agent(mock_context, request)
+
+        # Check that system role was stored
+        history = entity.state.conversation_history
+        assert _role_value(history[0]) == "system"
+        assert history[0].text == "System message"
+
+    async def test_run_agent_with_response_format(self) -> None:
+        """Test run_agent with a JSON response format."""
+        mock_agent = Mock()
+        # Return JSON response
+        mock_agent.run = AsyncMock(return_value=_agent_response('{"answer": 42}'))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request = RunRequest(
+            message="What is the answer?",
+            thread_id="conv-runreq-4",
+            response_format=EntityStructuredResponse,
+            correlation_id="corr-runreq-4",
+        )
+
+        result = await entity.run_agent(mock_context, request)
+
+        assert result["status"] == "success"
+        # Should have structured_response
+        if "structured_response" in result:
+            assert result["structured_response"]["answer"] == 42
+
+    async def test_run_agent_disable_tool_calls(self) -> None:
+        """Test run_agent with tool calls disabled."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity = AgentEntity(mock_agent)
+        mock_context = Mock()
+
+        request = RunRequest(
+            message="Test", thread_id="conv-runreq-5", enable_tool_calls=False, correlation_id="corr-runreq-5"
+        )
+
+        result = await entity.run_agent(mock_context, request)
+
+        assert result["status"] == "success"
+        # Agent should have been called (tool disabling is framework-dependent)
+        mock_agent.run.assert_called_once()
+
+    async def test_entity_function_with_run_request_dict(self) -> None:
+        """Test that the entity function handles the RunRequest dict format."""
+        mock_agent = Mock()
+        mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
+
+        entity_function = create_agent_entity(mock_agent)
+
+        mock_context = Mock()
+        mock_context.operation_name = "run_agent"
+        mock_context.get_input.return_value = {
+            "message": "Test message",
+            "thread_id": "conv-789",
+            "role": "user",
+            "enable_tool_calls": True,
+            "correlation_id": "corr-runreq-6",
+        }
+        mock_context.get_state.return_value = None
+
+        await asyncio.to_thread(entity_function, mock_context)
+
+        # Verify result was set
+        assert mock_context.set_result.called
+        result = mock_context.set_result.call_args[0][0]
+        assert result["status"] == "success"
+        assert result["message"] == "Test message"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])
@@ -0,0 +1,470 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for data models (AgentSessionId, RunRequest, AgentResponse)."""
+
+import azure.durable_functions as df
+import pytest
+from agent_framework import Role
+from pydantic import BaseModel
+
+from agent_framework_azurefunctions._models import AgentResponse, AgentSessionId, RunRequest
+
+
+class ModuleStructuredResponse(BaseModel):
+    value: int
+
+
+class TestAgentSessionId:
+    """Test suite for AgentSessionId."""
+
+    def test_init_creates_session_id(self) -> None:
+        """Test that AgentSessionId initializes correctly."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key-123"
+
+    def test_with_random_key_generates_guid(self) -> None:
+        """Test that with_random_key generates a GUID."""
+        session_id = AgentSessionId.with_random_key(name="AgentEntity")
+
+        assert session_id.name == "AgentEntity"
+        assert len(session_id.key) == 32  # UUID hex is 32 chars
+        # Verify it's a valid hex string
+        int(session_id.key, 16)
+
+    def test_with_random_key_unique_keys(self) -> None:
+        """Test that with_random_key generates unique keys."""
+        session_id1 = AgentSessionId.with_random_key(name="AgentEntity")
+        session_id2 = AgentSessionId.with_random_key(name="AgentEntity")
+
+        assert session_id1.key != session_id2.key
+
+    def test_to_entity_id_conversion(self) -> None:
+        """Test conversion to EntityId."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key")
+        entity_id = session_id.to_entity_id()
+
+        assert isinstance(entity_id, df.EntityId)
+        assert entity_id.name == "dafx-AgentEntity"
+        assert entity_id.key == "test-key"
+
+    def test_from_entity_id_conversion(self) -> None:
+        """Test creation from EntityId."""
+        entity_id = df.EntityId(name="dafx-AgentEntity", key="test-key")
+        session_id = AgentSessionId.from_entity_id(entity_id)
+
+        assert isinstance(session_id, AgentSessionId)
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key"
+
+    def test_round_trip_entity_id_conversion(self) -> None:
+        """Test round-trip conversion to and from EntityId."""
+        original = AgentSessionId(name="AgentEntity", key="test-key")
+        entity_id = original.to_entity_id()
+        restored = AgentSessionId.from_entity_id(entity_id)
+
+        assert restored.name == original.name
+        assert restored.key == original.key
+
+    def test_str_representation(self) -> None:
+        """Test string representation."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
+        str_repr = str(session_id)
+
+        assert str_repr == "@AgentEntity@test-key-123"
+
+    def test_repr_representation(self) -> None:
+        """Test repr representation."""
+        session_id = AgentSessionId(name="AgentEntity", key="test-key")
+        repr_str = repr(session_id)
+
+        assert "AgentSessionId" in repr_str
+        assert "AgentEntity" in repr_str
+        assert "test-key" in repr_str
+
+    def test_parse_valid_session_id(self) -> None:
+        """Test parsing valid session ID string."""
+        session_id = AgentSessionId.parse("@AgentEntity@test-key-123")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "test-key-123"
+
+    def test_parse_invalid_format_no_prefix(self) -> None:
+        """Test parsing invalid format without @ prefix."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("AgentEntity@test-key")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+    def test_parse_invalid_format_single_part(self) -> None:
+        """Test parsing invalid format with single part."""
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.parse("@AgentEntity")
+
+        assert "Invalid agent session ID format" in str(exc_info.value)
+
+    def test_parse_with_multiple_at_signs_in_key(self) -> None:
+        """Test parsing with @ signs in the key."""
+        session_id = AgentSessionId.parse("@AgentEntity@key-with@symbols")
+
+        assert session_id.name == "AgentEntity"
+        assert session_id.key == "key-with@symbols"
+
+    def test_parse_round_trip(self) -> None:
+        """Test round-trip parse and string conversion."""
+        original = AgentSessionId(name="AgentEntity", key="test-key")
+        str_repr = str(original)
+        parsed = AgentSessionId.parse(str_repr)
+
+        assert parsed.name == original.name
+        assert parsed.key == original.key
+
+    def test_to_entity_name_adds_prefix(self) -> None:
+        """Test that to_entity_name adds the dafx- prefix."""
+        entity_name = AgentSessionId.to_entity_name("TestAgent")
+        assert entity_name == "dafx-TestAgent"
+
+    def test_from_entity_id_strips_prefix(self) -> None:
+        """Test that from_entity_id strips the dafx- prefix."""
+        entity_id = df.EntityId(name="dafx-TestAgent", key="key123")
+        session_id = AgentSessionId.from_entity_id(entity_id)
+
+        assert session_id.name == "TestAgent"
+        assert session_id.key == "key123"
+
+    def test_from_entity_id_raises_without_prefix(self) -> None:
+        """Test that from_entity_id raises ValueError when entity name lacks the prefix."""
+        entity_id = df.EntityId(name="TestAgent", key="key123")
+
+        with pytest.raises(ValueError) as exc_info:
+            AgentSessionId.from_entity_id(entity_id)
+
+        assert "not a valid agent session ID" in str(exc_info.value)
+        assert "dafx-" in str(exc_info.value)
+
+
+class TestRunRequest:
+    """Test suite for RunRequest."""
+
+    def test_init_with_defaults(self) -> None:
+        """Test RunRequest initialization with defaults."""
+        request = RunRequest(message="Hello", thread_id="thread-default")
+
+        assert request.message == "Hello"
+        assert request.role == Role.USER
+        assert request.response_format is None
+        assert request.enable_tool_calls is True
+        assert request.thread_id == "thread-default"
+
+    def test_init_with_all_fields(self) -> None:
+        """Test RunRequest initialization with all fields."""
+        schema = ModuleStructuredResponse
+        request = RunRequest(
+            message="Hello",
+            thread_id="thread-123",
+            role=Role.SYSTEM,
+            response_format=schema,
+            enable_tool_calls=False,
+        )
+
+        assert request.message == "Hello"
+        assert request.role == Role.SYSTEM
+        assert request.response_format is schema
+        assert request.enable_tool_calls is False
+        assert request.thread_id == "thread-123"
+
+    def test_init_coerces_string_role(self) -> None:
+        """Ensure string role values are coerced into Role instances."""
+        request = RunRequest(message="Hello", thread_id="thread-str-role", role="system")  # type: ignore[arg-type]
+
+        assert request.role == Role.SYSTEM
+
+    def test_to_dict_with_defaults(self) -> None:
+        """Test to_dict with default values."""
+        request = RunRequest(message="Test message", thread_id="thread-to-dict")
+        data = request.to_dict()
+
+        assert data["message"] == "Test message"
+        assert data["enable_tool_calls"] is True
+        assert data["role"] == "user"
+        assert "response_format" not in data or data["response_format"] is None
+        assert data["thread_id"] == "thread-to-dict"
+
+    def test_to_dict_with_all_fields(self) -> None:
+        """Test to_dict with all fields."""
+        schema = ModuleStructuredResponse
+        request = RunRequest(
+            message="Hello",
+            thread_id="thread-456",
+            role=Role.ASSISTANT,
+            response_format=schema,
+            enable_tool_calls=False,
+        )
+        data = request.to_dict()
+
+        assert data["message"] == "Hello"
+        assert data["role"] == "assistant"
+        assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
+        assert data["response_format"]["module"] == schema.__module__
+        assert data["response_format"]["qualname"] == schema.__qualname__
+        assert data["enable_tool_calls"] is False
+        assert data["thread_id"] == "thread-456"
+
+    def test_from_dict_with_defaults(self) -> None:
+        """Test from_dict with minimal data."""
+        data = {"message": "Hello", "thread_id": "thread-from-dict"}
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Hello"
+        assert request.role == Role.USER
+        assert request.enable_tool_calls is True
+        assert request.thread_id == "thread-from-dict"
+
+    def test_from_dict_with_all_fields(self) -> None:
+        """Test from_dict with all fields."""
+        data = {
+            "message": "Test",
+            "role": "system",
+            "response_format": {
+                "__response_schema_type__": "pydantic_model",
+                "module": ModuleStructuredResponse.__module__,
+                "qualname": ModuleStructuredResponse.__qualname__,
+            },
+            "enable_tool_calls": False,
+            "thread_id": "thread-789",
+        }
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Test"
+        assert request.role == Role.SYSTEM
+        assert request.response_format is ModuleStructuredResponse
+        assert request.enable_tool_calls is False
+        assert request.thread_id == "thread-789"
+
+    def test_from_dict_with_unknown_role_preserves_value(self) -> None:
+        """Test from_dict keeps custom roles intact."""
+        data = {"message": "Test", "role": "reviewer", "thread_id": "thread-with-custom-role"}
+        request = RunRequest.from_dict(data)
+
+        assert request.role.value == "reviewer"
+        assert request.role != Role.USER
+
+    def test_from_dict_empty_message(self) -> None:
+        """Test from_dict with empty message."""
+        data = {"thread_id": "thread-empty"}
+        request = RunRequest.from_dict(data)
+
+        assert request.message == ""
+        assert request.role == Role.USER
+        assert request.thread_id == "thread-empty"
+
+    def test_round_trip_dict_conversion(self) -> None:
+        """Test round-trip to_dict and from_dict."""
+        original = RunRequest(
+            message="Test message",
+            thread_id="thread-123",
+            role=Role.SYSTEM,
+            response_format=ModuleStructuredResponse,
+            enable_tool_calls=False,
+        )
+
+        data = original.to_dict()
+        restored = RunRequest.from_dict(data)
+
+        assert restored.message == original.message
+        assert restored.role == original.role
+        assert restored.response_format is ModuleStructuredResponse
+        assert restored.enable_tool_calls == original.enable_tool_calls
+        assert restored.thread_id == original.thread_id
+
+    def test_round_trip_with_pydantic_response_format(self) -> None:
+        """Ensure Pydantic response formats serialize and deserialize properly."""
+        original = RunRequest(
+            message="Structured",
+            thread_id="thread-pydantic",
+            response_format=ModuleStructuredResponse,
+        )
+
+        data = original.to_dict()
+
+        assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
+        assert data["response_format"]["module"] == ModuleStructuredResponse.__module__
+        assert data["response_format"]["qualname"] == ModuleStructuredResponse.__qualname__
+
+        restored = RunRequest.from_dict(data)
+        assert restored.response_format is ModuleStructuredResponse
+
+    def test_init_with_correlation_id(self) -> None:
+        """Test RunRequest initialization with correlation_id."""
+        request = RunRequest(message="Test message", thread_id="thread-corr-init", correlation_id="corr-123")
+
+        assert request.message == "Test message"
+        assert request.correlation_id == "corr-123"
+
+    def test_to_dict_with_correlation_id(self) -> None:
+        """Test to_dict includes correlation_id."""
+        request = RunRequest(message="Test", thread_id="thread-corr-to-dict", correlation_id="corr-456")
+        data = request.to_dict()
+
+        assert data["message"] == "Test"
+        assert data["correlation_id"] == "corr-456"
+
+    def test_from_dict_with_correlation_id(self) -> None:
+        """Test from_dict with correlation_id."""
+        data = {"message": "Test", "correlation_id": "corr-789", "thread_id": "thread-corr-from-dict"}
+        request = RunRequest.from_dict(data)
+
+        assert request.message == "Test"
+        assert request.correlation_id == "corr-789"
+        assert request.thread_id == "thread-corr-from-dict"
+
+    def test_round_trip_with_correlation_id(self) -> None:
+        """Test round-trip to_dict and from_dict with correlation_id."""
+        original = RunRequest(
+            message="Test message",
+            thread_id="thread-123",
+            role=Role.SYSTEM,
+            correlation_id="corr-123",
+        )
+
+        data = original.to_dict()
+        restored = RunRequest.from_dict(data)
+
+        assert restored.message == original.message
+        assert restored.role == original.role
+        assert restored.correlation_id == original.correlation_id
+        assert restored.thread_id == original.thread_id
+
+
+class TestAgentResponse:
+    """Test suite for AgentResponse."""
+
+    def test_init_with_required_fields(self) -> None:
+        """Test AgentResponse initialization with required fields."""
+        response = AgentResponse(
+            response="Test response", message="Test message", thread_id="thread-123", status="success"
+        )
+
+        assert response.response == "Test response"
+        assert response.message == "Test message"
+        assert response.thread_id == "thread-123"
+        assert response.status == "success"
+        assert response.message_count == 0
+        assert response.error is None
+        assert response.error_type is None
+        assert response.structured_response is None
+
+    def test_init_with_all_fields(self) -> None:
+        """Test AgentResponse initialization with all fields."""
+        structured = {"answer": "42"}
+        response = AgentResponse(
+            response=None,
+            message="What is the answer?",
+            thread_id="thread-456",
+            status="success",
+            message_count=5,
+            error=None,
+            error_type=None,
+            structured_response=structured,
+        )
+
+        assert response.response is None
+        assert response.structured_response == structured
+        assert response.message_count == 5
+
+    def test_to_dict_with_text_response(self) -> None:
+        """Test to_dict with text response."""
+        response = AgentResponse(
+            response="Text response", message="Message", thread_id="thread-1", status="success", message_count=3
+        )
+        data = response.to_dict()
+
+        assert data["response"] == "Text response"
+        assert data["message"] == "Message"
+        assert data["thread_id"] == "thread-1"
+        assert data["status"] == "success"
+        assert data["message_count"] == 3
+        assert "structured_response" not in data
+        assert "error" not in data
+        assert "error_type" not in data
+
+    def test_to_dict_with_structured_response(self) -> None:
+        """Test to_dict with structured response."""
+        structured = {"answer": 42, "confidence": 0.95}
+        response = AgentResponse(
+            response=None,
+            message="Question",
+            thread_id="thread-2",
+            status="success",
+            structured_response=structured,
+        )
+        data = response.to_dict()
+
+        assert data["structured_response"] == structured
+        assert "response" not in data
+
+    def test_to_dict_with_error(self) -> None:
+        """Test to_dict with error."""
+        response = AgentResponse(
+            response=None,
+            message="Failed message",
+            thread_id="thread-3",
+            status="error",
+            error="Something went wrong",
+            error_type="ValueError",
+        )
+        data = response.to_dict()
+
+        assert data["status"] == "error"
+        assert data["error"] == "Something went wrong"
+        assert data["error_type"] == "ValueError"
+
+    def test_to_dict_prefers_structured_over_text(self) -> None:
+        """Test to_dict prefers structured_response over response."""
+        structured = {"result": "structured"}
+        response = AgentResponse(
+            response="Text response",
+            message="Message",
+            thread_id="thread-4",
+            status="success",
+            structured_response=structured,
+        )
+        data = response.to_dict()
+
+        assert "structured_response" in data
+        assert data["structured_response"] == structured
+        # Text response should not be included when structured is present
+        assert "response" not in data
+
+
+class TestModelIntegration:
+    """Test suite for integration between models."""
+
+    def test_run_request_with_session_id(self) -> None:
+        """Test using RunRequest with AgentSessionId."""
+        session_id = AgentSessionId.with_random_key("AgentEntity")
+        request = RunRequest(message="Test message", thread_id=str(session_id))
+
+        assert request.thread_id is not None
+        assert request.thread_id == str(session_id)
+        assert request.thread_id.startswith("@AgentEntity@")
+
+    def test_response_from_run_request(self) -> None:
+        """Test creating AgentResponse from RunRequest."""
+        request = RunRequest(message="What is 2+2?", thread_id="thread-123", role=Role.USER)
+
+        response = AgentResponse(
+            response="4",
+            message=request.message,
+            thread_id=request.thread_id,
+            status="success",
+            message_count=1,
+        )
+
+        assert response.message == request.message
+        assert response.thread_id == request.thread_id
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])
@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for multi-agent support in AgentFunctionApp."""
+
+from unittest.mock import Mock
+
+import pytest
+
+from agent_framework_azurefunctions import AgentFunctionApp
+
+
+class TestMultiAgentInit:
+    """Test suite for multi-agent initialization."""
+
+    def test_init_with_agents_list(self) -> None:
+        """Test initialization with list of agents."""
+        agent1 = Mock()
+        agent1.name = "Agent1"
+        agent2 = Mock()
+        agent2.name = "Agent2"
+
+        app = AgentFunctionApp(agents=[agent1, agent2])
+
+        assert len(app.agents) == 2
+        assert "Agent1" in app.agents
+        assert "Agent2" in app.agents
+        assert app.agents["Agent1"] == agent1
+        assert app.agents["Agent2"] == agent2
+
+    def test_init_with_empty_agents_list(self) -> None:
+        """Test initialization with empty list of agents."""
+        app = AgentFunctionApp(agents=[])
+
+        assert len(app.agents) == 0
+
+    def test_init_with_no_agents(self) -> None:
+        """Test initialization without any agents."""
+        app = AgentFunctionApp()
+
+        assert len(app.agents) == 0
+
+    def test_init_with_duplicate_agent_names(self) -> None:
+        """Test initialization with agents having the same name raises error."""
+        agent1 = Mock()
+        agent1.name = "TestAgent"
+        agent2 = Mock()
+        agent2.name = "TestAgent"
+
+        with pytest.raises(ValueError, match="already registered"):
+            AgentFunctionApp(agents=[agent1, agent2])
+
+    def test_init_with_agent_without_name(self) -> None:
+        """Test initialization with agent missing name attribute raises error."""
+        agent1 = Mock()
+        agent1.name = "Agent1"
+        agent2 = Mock(spec=[])  # Mock without name attribute
+
+        with pytest.raises(ValueError, match="does not have a 'name' attribute"):
+            AgentFunctionApp(agents=[agent1, agent2])
+
+
+class TestAddAgentMethod:
+    """Test suite for add_agent() method."""
+
+    def test_add_agent_to_empty_app(self) -> None:
+        """Test adding agent to app initialized without agents."""
+        app = AgentFunctionApp()
+
+        agent = Mock()
+        agent.name = "NewAgent"
+
+        app.add_agent(agent)
+
+        assert len(app.agents) == 1
+        assert "NewAgent" in app.agents
+        assert app.agents["NewAgent"] == agent
+
+    def test_add_multiple_agents(self) -> None:
+        """Test adding multiple agents sequentially."""
+        app = AgentFunctionApp()
+
+        agent1 = Mock()
+        agent1.name = "Agent1"
+        agent2 = Mock()
+        agent2.name = "Agent2"
+
+        app.add_agent(agent1)
+        app.add_agent(agent2)
+
+        assert len(app.agents) == 2
+        assert "Agent1" in app.agents
+        assert "Agent2" in app.agents
+
+    def test_add_agent_with_duplicate_name_raises_error(self) -> None:
+        """Test that adding agent with duplicate name raises ValueError."""
+        agent1 = Mock()
+        agent1.name = "MyAgent"
+        agent2 = Mock()
+        agent2.name = "MyAgent"
+
+        app = AgentFunctionApp(agents=[agent1])
+
+        # Try to add another agent with the same name
+        with pytest.raises(ValueError, match="already registered"):
+            app.add_agent(agent2)
+
+    def test_add_agent_to_app_with_existing_agents(self) -> None:
+        """Test adding agent to app that already has agents."""
+        agent1 = Mock()
+        agent1.name = "Agent1"
+        agent2 = Mock()
+        agent2.name = "Agent2"
+
+        app = AgentFunctionApp(agents=[agent1])
+        app.add_agent(agent2)
+
+        assert len(app.agents) == 2
+        assert "Agent1" in app.agents
+        assert "Agent2" in app.agents
+
+    def test_add_agent_without_name_raises_error(self) -> None:
+        """Test that adding agent without name attribute raises error."""
+        app = AgentFunctionApp()
+
+        agent = Mock(spec=[])  # Mock without name attribute
+
+        with pytest.raises(ValueError, match="does not have a 'name' attribute"):
+            app.add_agent(agent)
+
+
+class TestHealthCheckWithMultipleAgents:
+    """Test suite for health check with multiple agents."""
+
+    def test_health_check_returns_all_agents(self) -> None:
+        """Test that health check returns information about all agents."""
+        agent1 = Mock()
+        agent1.name = "Agent1"
+        agent2 = Mock()
+        agent2.name = "Agent2"
+
+        app = AgentFunctionApp(agents=[agent1, agent2])
+
+        # Note: We can't easily test the actual health check endpoint without running the app
+        # But we can verify the agents dictionary is properly populated
+        assert len(app.agents) == 2
+        assert app.enable_health_check is True
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])
@@ -0,0 +1,442 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for orchestration support (DurableAIAgent)."""
+
+from typing import Any
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import AgentThread
+
+from agent_framework_azurefunctions import AgentFunctionApp, DurableAIAgent
+from agent_framework_azurefunctions._models import AgentSessionId, DurableAgentThread
+
+
+def _app_with_registered_agents(*agent_names: str) -> AgentFunctionApp:
+    app = AgentFunctionApp(enable_health_check=False, enable_http_endpoints=False)
+    for name in agent_names:
+        agent = Mock()
+        agent.name = name
+        app.add_agent(agent)
+    return app
+
+
+class TestDurableAIAgent:
+    """Test suite for DurableAIAgent wrapper."""
+
+    def test_init(self) -> None:
+        """Test DurableAIAgent initialization."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-123"
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        assert agent.context == mock_context
+        assert agent.agent_name == "TestAgent"
+
+    def test_implements_agent_protocol(self) -> None:
+        """Test that DurableAIAgent implements AgentProtocol."""
+        from agent_framework import AgentProtocol
+
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # Check that agent satisfies AgentProtocol
+        assert isinstance(agent, AgentProtocol)
+
+    def test_has_agent_protocol_properties(self) -> None:
+        """Test that DurableAIAgent has AgentProtocol properties."""
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # AgentProtocol properties
+        assert hasattr(agent, "id")
+        assert hasattr(agent, "name")
+        assert hasattr(agent, "description")
+        assert hasattr(agent, "display_name")
+
+        # Verify values
+        assert agent.name == "TestAgent"
+        assert agent.description == "Durable agent proxy for TestAgent"
+        assert agent.display_name == "TestAgent"
+        assert agent.id is not None  # Auto-generated UUID
+
+    def test_get_new_thread(self) -> None:
+        """Test creating a new agent thread."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-456"
+        mock_context.new_uuid = Mock(return_value="test-guid-456")
+
+        agent = DurableAIAgent(mock_context, "WriterAgent")
+        thread = agent.get_new_thread()
+
+        assert isinstance(thread, DurableAgentThread)
+        assert thread.session_id is not None
+        session_id = thread.session_id
+        assert isinstance(session_id, AgentSessionId)
+        assert session_id.name == "WriterAgent"
+        assert session_id.key == "test-guid-456"
+        mock_context.new_uuid.assert_called_once()
+
+    def test_get_new_thread_deterministic(self) -> None:
+        """Test that get_new_thread creates deterministic session IDs."""
+
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-789"
+        mock_context.new_uuid = Mock(side_effect=["session-guid-1", "session-guid-2"])
+
+        agent = DurableAIAgent(mock_context, "EditorAgent")
+
+        # Create multiple threads - they should have unique session IDs
+        thread1 = agent.get_new_thread()
+        thread2 = agent.get_new_thread()
+
+        assert isinstance(thread1, DurableAgentThread)
+        assert isinstance(thread2, DurableAgentThread)
+
+        session_id1 = thread1.session_id
+        session_id2 = thread2.session_id
+        assert session_id1 is not None and session_id2 is not None
+        assert isinstance(session_id1, AgentSessionId)
+        assert isinstance(session_id2, AgentSessionId)
+        assert session_id1.name == "EditorAgent"
+        assert session_id2.name == "EditorAgent"
+        assert session_id1.key == "session-guid-1"
+        assert session_id2.key == "session-guid-2"
+        assert mock_context.new_uuid.call_count == 2
+
+    def test_run_creates_entity_call(self) -> None:
+        """Test that run() creates proper entity call and returns a Task."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-001"
+        mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
+
+        # Mock call_entity to return a Task-like object
+        mock_task = Mock()
+        mock_task._is_scheduled = False  # Task attribute that orchestration checks
+
+        mock_context.call_entity = Mock(return_value=mock_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # Create thread
+        thread = agent.get_new_thread()
+
+        # Call run() - it should return the Task directly
+        task = agent.run(messages="Test message", thread=thread, enable_tool_calls=True)
+
+        # Verify run() returns the Task from call_entity
+        assert task == mock_task
+
+        # Verify call_entity was called with correct parameters
+        assert mock_context.call_entity.called
+        call_args = mock_context.call_entity.call_args
+        entity_id, operation, request = call_args[0]
+
+        assert operation == "run_agent"
+        assert request["message"] == "Test message"
+        assert request["enable_tool_calls"] is True
+        assert "correlation_id" in request
+        assert request["correlation_id"] == "correlation-guid"
+        assert "thread_id" in request
+        assert request["thread_id"] == "thread-guid"
+
+    def test_run_without_thread(self) -> None:
+        """Test that run() works without explicit thread (creates unique session key)."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-002"
+        # Two calls to new_uuid: one for session_key, one for correlation_id
+        mock_context.new_uuid = Mock(side_effect=["auto-generated-guid", "correlation-guid"])
+
+        mock_task = Mock()
+        mock_task._is_scheduled = False
+        mock_context.call_entity = Mock(return_value=mock_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        # Call without thread
+        task = agent.run(messages="Test message")
+
+        assert task == mock_task
+
+        # Verify the entity ID uses the auto-generated GUID with dafx- prefix
+        call_args = mock_context.call_entity.call_args
+        entity_id = call_args[0][0]
+        assert entity_id.name == "dafx-TestAgent"
+        assert entity_id.key == "auto-generated-guid"
+        # Should be called twice: once for session_key, once for correlation_id
+        assert mock_context.new_uuid.call_count == 2
+
+    def test_run_with_response_format(self) -> None:
+        """Test that run() passes response format correctly."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-003"
+
+        mock_task = Mock()
+        mock_task._is_scheduled = False
+        mock_context.call_entity = Mock(return_value=mock_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        from pydantic import BaseModel
+
+        class SampleSchema(BaseModel):
+            key: str
+
+        # Create thread and call
+        thread = agent.get_new_thread()
+
+        task = agent.run(messages="Test message", thread=thread, response_format=SampleSchema)
+
+        assert task == mock_task
+
+        # Verify schema was passed in the call_entity arguments
+        call_args = mock_context.call_entity.call_args
+        input_data = call_args[0][2]  # Third argument is input_data
+        assert "response_format" in input_data
+        assert input_data["response_format"]["__response_schema_type__"] == "pydantic_model"
+        assert input_data["response_format"]["module"] == SampleSchema.__module__
+        assert input_data["response_format"]["qualname"] == SampleSchema.__qualname__
+
+    def test_messages_to_string(self) -> None:
+        """Test converting ChatMessage list to string."""
+        from agent_framework import ChatMessage
+
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        messages = [
+            ChatMessage(role="user", text="Hello"),
+            ChatMessage(role="assistant", text="Hi there"),
+            ChatMessage(role="user", text="How are you?"),
+        ]
+
+        result = agent._messages_to_string(messages)
+
+        assert result == "Hello\nHi there\nHow are you?"
+
+    def test_run_with_chat_message(self) -> None:
+        """Test that run() handles ChatMessage input."""
+        from agent_framework import ChatMessage
+
+        mock_context = Mock()
+        mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
+        mock_task = Mock()
+        mock_context.call_entity = Mock(return_value=mock_task)
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+        thread = agent.get_new_thread()
+
+        # Call with ChatMessage
+        msg = ChatMessage(role="user", text="Hello")
+        task = agent.run(messages=msg, thread=thread)
+
+        assert task == mock_task
+
+        # Verify message was converted to string
+        call_args = mock_context.call_entity.call_args
+        request = call_args[0][2]
+        assert request["message"] == "Hello"
+
+    def test_run_stream_raises_not_implemented(self) -> None:
+        """Test that run_stream() method raises NotImplementedError."""
+        mock_context = Mock()
+        agent = DurableAIAgent(mock_context, "TestAgent")
+
+        with pytest.raises(NotImplementedError) as exc_info:
+            agent.run_stream("Test message")
+
+        error_msg = str(exc_info.value)
+        assert "Streaming is not supported" in error_msg
+
+    def test_entity_id_format(self) -> None:
+        """Test that EntityId is created with correct format (name, key)."""
+        from azure.durable_functions import EntityId
+
+        mock_context = Mock()
+        mock_context.new_uuid = Mock(return_value="test-guid-789")
+        mock_context.call_entity = Mock(return_value=Mock())
+
+        agent = DurableAIAgent(mock_context, "WriterAgent")
+        thread = agent.get_new_thread()
+
+        # Call run() to trigger entity ID creation
+        agent.run("Test", thread=thread)
+
+        # Verify call_entity was called with correct EntityId
+        call_args = mock_context.call_entity.call_args
+        entity_id = call_args[0][0]
+
+        # EntityId should be EntityId(name="dafx-WriterAgent", key="test-guid-789")
+        # Which formats as "@dafx-writeragent@test-guid-789"
+        assert isinstance(entity_id, EntityId)
+        assert entity_id.name == "dafx-WriterAgent"
+        assert entity_id.key == "test-guid-789"
+        assert str(entity_id) == "@dafx-writeragent@test-guid-789"
+
+
+class TestAgentFunctionAppGetAgent:
+    """Test suite for AgentFunctionApp.get_agent."""
+
+    def test_get_agent_method(self) -> None:
+        """Test get_agent method creates DurableAIAgent for registered agent."""
+        app = _app_with_registered_agents("MyAgent")
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-100"
+
+        agent = app.get_agent(mock_context, "MyAgent")
+
+        assert isinstance(agent, DurableAIAgent)
+        assert agent.agent_name == "MyAgent"
+        assert agent.context == mock_context
+
+    def test_get_agent_raises_for_unregistered_agent(self) -> None:
+        """Test get_agent raises ValueError when agent is not registered."""
+        app = _app_with_registered_agents("KnownAgent")
+
+        with pytest.raises(ValueError, match=r"Agent 'MissingAgent' is not registered with this app\."):
+            app.get_agent(Mock(), "MissingAgent")
+
+
+class TestOrchestrationIntegration:
+    """Integration tests for orchestration scenarios."""
+
+    def test_sequential_agent_calls_simulation(self) -> None:
+        """Simulate sequential agent calls in an orchestration."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-orchestration-001"
+        # new_uuid will be called 3 times:
+        # 1. thread creation
+        # 2. correlation_id for first call
+        # 3. correlation_id for second call
+        mock_context.new_uuid = Mock(side_effect=["deterministic-guid-001", "corr-1", "corr-2"])
+
+        # Track entity calls
+        entity_calls: list[dict[str, Any]] = []
+
+        def mock_call_entity_side_effect(entity_id: Any, operation: str, input_data: dict[str, Any]) -> Mock:
+            entity_calls.append({"entity_id": str(entity_id), "operation": operation, "input": input_data})
+
+            # Return a mock Task
+            mock_task = Mock()
+            mock_task._is_scheduled = False
+            return mock_task
+
+        mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
+
+        app = _app_with_registered_agents("WriterAgent")
+        agent = app.get_agent(mock_context, "WriterAgent")
+
+        # Create thread
+        thread = agent.get_new_thread()
+
+        # First call - returns Task
+        task1 = agent.run("Write something", thread=thread)
+        assert hasattr(task1, "_is_scheduled")
+
+        # Second call - returns Task
+        task2 = agent.run("Improve: something", thread=thread)
+        assert hasattr(task2, "_is_scheduled")
+
+        # Verify both calls used the same entity (same session key)
+        assert len(entity_calls) == 2
+        assert entity_calls[0]["entity_id"] == entity_calls[1]["entity_id"]
+        # EntityId format is @dafx-writeragent@deterministic-guid-001
+        assert entity_calls[0]["entity_id"] == "@dafx-writeragent@deterministic-guid-001"
+        # new_uuid called 3 times: thread + 2 correlation IDs
+        assert mock_context.new_uuid.call_count == 3
+
+    def test_multiple_agents_in_orchestration(self) -> None:
+        """Test using multiple different agents in one orchestration."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-orchestration-002"
+        # Mock new_uuid to return different GUIDs for each call
+        # Order: writer thread, editor thread, writer correlation, editor correlation
+        mock_context.new_uuid = Mock(side_effect=["writer-guid-001", "editor-guid-002", "writer-corr", "editor-corr"])
+
+        entity_calls: list[str] = []
+
+        def mock_call_entity_side_effect(entity_id: Any, operation: str, input_data: dict[str, Any]) -> Mock:
+            entity_calls.append(str(entity_id))
+            mock_task = Mock()
+            mock_task._is_scheduled = False
+            return mock_task
+
+        mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
+
+        app = _app_with_registered_agents("WriterAgent", "EditorAgent")
+        writer = app.get_agent(mock_context, "WriterAgent")
+        editor = app.get_agent(mock_context, "EditorAgent")
+
+        writer_thread = writer.get_new_thread()
+        editor_thread = editor.get_new_thread()
+
+        # Call both agents - returns Tasks
+        writer_task = writer.run("Write", thread=writer_thread)
+        editor_task = editor.run("Edit", thread=editor_thread)
+
+        assert hasattr(writer_task, "_is_scheduled")
+        assert hasattr(editor_task, "_is_scheduled")
+
+        # Verify different entity IDs were used
+        assert len(entity_calls) == 2
+        # EntityId format is @dafx-agentname@guid (lowercased agent name with dafx- prefix)
+        assert entity_calls[0] == "@dafx-writeragent@writer-guid-001"
+        assert entity_calls[1] == "@dafx-editoragent@editor-guid-002"
+
+
+class TestAgentThreadSerialization:
+    """Test that AgentThread can be serialized for orchestration state."""
+
+    async def test_agent_thread_serialize(self) -> None:
+        """Test that AgentThread can be serialized."""
+        thread = AgentThread()
+
+        # Serialize
+        serialized = await thread.serialize()
+
+        assert isinstance(serialized, dict)
+        assert "service_thread_id" in serialized
+
+    async def test_agent_thread_deserialize(self) -> None:
+        """Test that AgentThread can be deserialized."""
+        thread = AgentThread()
+        serialized = await thread.serialize()
+
+        # Deserialize
+        restored = await AgentThread.deserialize(serialized)
+
+        assert isinstance(restored, AgentThread)
+        assert restored.service_thread_id == thread.service_thread_id
+
+    async def test_durable_agent_thread_serialization(self) -> None:
+        """Test that DurableAgentThread persists session metadata during serialization."""
+        mock_context = Mock()
+        mock_context.instance_id = "test-instance-999"
+        mock_context.new_uuid = Mock(return_value="test-guid-999")
+
+        agent = DurableAIAgent(mock_context, "TestAgent")
+        thread = agent.get_new_thread()
+
+        assert isinstance(thread, DurableAgentThread)
+        # Verify custom attribute and property exist
+        assert thread.session_id is not None
+        session_id = thread.session_id
+        assert isinstance(session_id, AgentSessionId)
+        assert session_id.name == "TestAgent"
+        assert session_id.key == "test-guid-999"
+
+        # Standard serialization should still work
+        serialized = await thread.serialize()
+        assert isinstance(serialized, dict)
+        assert serialized.get("durable_session_id") == str(session_id)
+
+        # After deserialization, we'd need to restore the custom attribute
+        # This would be handled by the orchestration framework
+        restored = await DurableAgentThread.deserialize(serialized)
+        assert isinstance(restored, DurableAgentThread)
+        assert restored.session_id == session_id
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])
@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Unit tests for AgentState correlation ID tracking."""
+
+from unittest.mock import Mock
+
+import pytest
+from agent_framework import AgentRunResponse
+
+from agent_framework_azurefunctions._state import AgentState
+
+
+class TestAgentStateCorrelationId:
+    """Test suite for AgentState correlation ID tracking."""
+
+    def _create_mock_response(self, text: str = "Response") -> Mock:
+        """Create a mock AgentRunResponse with the provided text."""
+        mock_response = Mock(spec=AgentRunResponse)
+        mock_response.to_dict.return_value = {"text": text, "messages": []}
+        return mock_response
+
+    def test_add_assistant_message_with_correlation_id(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-123-request")
+        state.add_assistant_message("Response", self._create_mock_response(), correlation_id="corr-123")
+        message_metadata = state.conversation_history[-1].additional_properties or {}
+        assert message_metadata.get("correlation_id") == "corr-123"
+
+        response_data = state.try_get_agent_response("corr-123")
+        assert response_data is not None
+        assert response_data["content"] == "Response"
+        assert response_data["agent_response"] == {"text": "Response", "messages": []}
+
+    def test_try_get_agent_response_returns_response(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-200-request")
+        state.add_assistant_message("Response", self._create_mock_response(), correlation_id="corr-456")
+
+        response_data = state.try_get_agent_response("corr-456")
+
+        assert response_data is not None
+        assert response_data["content"] == "Response"
+
+    def test_try_get_agent_response_returns_none_for_missing_id(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-300-request")
+        state.add_assistant_message("Response", self._create_mock_response(), correlation_id="corr-123")
+
+        assert state.try_get_agent_response("non-existent") is None
+
+    def test_multiple_responses_tracked_separately(self) -> None:
+        state = AgentState()
+
+        for index in range(3):
+            state.add_user_message(f"Message {index}", correlation_id=f"corr-{index}-request")
+            state.add_assistant_message(
+                f"Response {index}",
+                self._create_mock_response(text=f"Response {index}"),
+                correlation_id=f"corr-{index}",
+            )
+
+        for index in range(3):
+            payload = state.try_get_agent_response(f"corr-{index}")
+            assert payload is not None
+            assert payload["content"] == f"Response {index}"
+
+    def test_add_assistant_message_without_correlation_id(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-400-request")
+        state.add_assistant_message("Response", self._create_mock_response())
+
+        assert state.try_get_agent_response("missing") is None
+        assert state.last_response == "Response"
+
+    def test_to_dict_does_not_duplicate_agent_responses(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-500-request")
+        state.add_assistant_message("Response", self._create_mock_response(), correlation_id="corr-123")
+
+        state_snapshot = state.to_dict()
+
+        assert "agent_responses" not in state_snapshot
+        metadata = state_snapshot["conversation_history"][-1]["additional_properties"]
+        assert metadata["correlation_id"] == "corr-123"
+
+    def test_restore_state_preserves_agent_response_lookup(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-600-request")
+        state.add_assistant_message("Response", self._create_mock_response(), correlation_id="corr-123")
+
+        restored_state = AgentState()
+        restored_state.restore_state(state.to_dict())
+
+        payload = restored_state.try_get_agent_response("corr-123")
+        assert payload is not None
+        assert payload["content"] == "Response"
+
+    def test_reset_clears_conversation_history(self) -> None:
+        state = AgentState()
+        state.add_user_message("Hello", correlation_id="corr-700-request")
+        state.add_assistant_message("Response", self._create_mock_response(), correlation_id="corr-123")
+
+        state.reset()
+
+        assert len(state.conversation_history) == 0
+        assert state.try_get_agent_response("corr-123") is None
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "--tb=short"])
				`@@ -0,0 +1 @@`
				`# Copyright (c) Microsoft. All rights reserved.`