agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Add Entity State Providers for DurableTask Package (#2981)

Browse Source 
* Add Entity State Providers

* address comments

* Fix tests

* Fix tests

* Revert unrelated changes and remove thread_id

* Revert unrelated files
This commit is contained in:
Laveesh Rohra
2025-12-22 12:54:24 -08:00
committed by GitHub
Unverified
parent 87a38bc7da
commit a02527f00a
16 changed files with 1276 additions and 1432 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azurefunctions/tests/test_app.py
+61 -48
View File
@@ -2,6 +2,8 @@
"""Unit tests for AgentFunctionApp."""
# pyright: reportPrivateUsage=false
import json
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar
@@ -17,15 +19,36 @@ from agent_framework_durabletask import (
THREAD_ID_HEADER,
WAIT_FOR_RESPONSE_FIELD,
WAIT_FOR_RESPONSE_HEADER,
AgentEntity,
AgentEntityStateProviderMixin,
DurableAgentState,
)
from agent_framework_azurefunctions import AgentFunctionApp
from agent_framework_azurefunctions._entities import AgentEntity, create_agent_entity
from agent_framework_azurefunctions._entities import create_agent_entity
TFunc = TypeVar("TFunc", bound=Callable[..., Any])
def _identity_decorator(func: TFunc) -> TFunc:
return func
class _InMemoryStateProvider(AgentEntityStateProviderMixin):
def __init__(self, *, thread_id: str = "test-thread", initial_state: dict[str, Any] | None = None) -> None:
self._thread_id = thread_id
self._state_dict: dict[str, Any] = initial_state or {}
def _get_state_dict(self) -> dict[str, Any]:
return self._state_dict
def _set_state_dict(self, state: dict[str, Any]) -> None:
self._state_dict = state
def _get_thread_id_from_entity(self) -> str:
return self._thread_id
class TestAgentFunctionAppInit:
"""Test suite for AgentFunctionApp initialization."""
@@ -89,7 +112,7 @@ class TestAgentFunctionAppInit:
app.add_agent(mock_agent, callback=specific_callback)
setup_mock.assert_called_once()
_, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
_, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
assert passed_callback is specific_callback
assert enable_http_endpoint is True
@@ -105,7 +128,7 @@ class TestAgentFunctionAppInit:
app.add_agent(mock_agent)
setup_mock.assert_called_once()
_, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
_, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
assert passed_callback is default_callback
assert enable_http_endpoint is True
@@ -120,7 +143,7 @@ class TestAgentFunctionAppInit:
AgentFunctionApp(agents=[mock_agent], default_callback=default_callback)
setup_mock.assert_called_once()
_, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
_, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
assert passed_callback is default_callback
assert enable_http_endpoint is True
@@ -336,13 +359,12 @@ class TestAgentEntityOperations:
return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Test response")])
)
entity = AgentEntity(mock_agent)
mock_context = Mock()
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="test-conv-123"))
result = await entity.run(
mock_context,
{"message": "Test message", "thread_id": "test-conv-123", "correlationId": "corr-app-entity-1"},
)
result = await entity.run({
"message": "Test message",
"correlationId": "corr-app-entity-1",
})
assert isinstance(result, AgentRunResponse)
assert result.text == "Test response"
@@ -355,22 +377,17 @@ class TestAgentEntityOperations:
return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response 1")])
)
entity = AgentEntity(mock_agent)
mock_context = Mock()
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
# Send first message
await entity.run(
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-app-entity-2"}
)
await entity.run({"message": "Message 1", "correlationId": "corr-app-entity-2"})
# Each conversation turn creates 2 entries: request and response
history = entity.state.data.conversation_history[0].messages # Request entry
assert len(history) == 1 # Just the user message
# Send second message
await entity.run(
mock_context, {"message": "Message 2", "thread_id": "conv-2", "correlationId": "corr-app-entity-2b"}
)
await entity.run({"message": "Message 2", "correlationId": "corr-app-entity-2b"})
# Now we have 4 entries total (2 requests + 2 responses)
# Access the first request entry
@@ -394,32 +411,26 @@ class TestAgentEntityOperations:
return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response")])
)
entity = AgentEntity(mock_agent)
mock_context = Mock()
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
assert len(entity.state.data.conversation_history) == 0
await entity.run(
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-app-entity-3a"}
)
await entity.run({"message": "Message 1", "correlationId": "corr-app-entity-3a"})
assert len(entity.state.data.conversation_history) == 2
await entity.run(
mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-app-entity-3b"}
)
await entity.run({"message": "Message 2", "correlationId": "corr-app-entity-3b"})
assert len(entity.state.data.conversation_history) == 4
def test_entity_reset(self) -> None:
"""Test that entity reset clears state."""
mock_agent = Mock()
entity = AgentEntity(mock_agent)
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider())
# Set some state
entity.state = DurableAgentState()
# Reset
mock_context = Mock()
entity.reset(mock_context)
entity.reset()
assert len(entity.state.data.conversation_history) == 0
@@ -448,7 +459,6 @@ class TestAgentEntityFactory:
mock_context.operation_name = "run"
mock_context.get_input.return_value = {
"message": "Test message",
"thread_id": "conv-123",
"correlationId": "corr-app-factory-1",
}
mock_context.get_state.return_value = None
@@ -476,7 +486,6 @@ class TestAgentEntityFactory:
mock_context.operation_name = "run_agent"
mock_context.get_input.return_value = {
"message": "Test message",
"thread_id": "conv-123",
"correlationId": "corr-app-factory-1",
}
mock_context.get_state.return_value = None
@@ -596,7 +605,11 @@ class TestAgentEntityFactory:
}
mock_context = Mock()
mock_context.operation_name = "reset"
mock_context.operation_name = "run"
mock_context.get_input.return_value = {
"message": "Test message",
"correlationId": "corr-restore-1",
}
mock_context.get_state.return_value = existing_state
with patch.object(DurableAgentState, "from_dict", wraps=DurableAgentState.from_dict) as from_dict_mock:
@@ -613,12 +626,12 @@ class TestErrorHandling:
mock_agent = Mock()
mock_agent.run = AsyncMock(side_effect=Exception("Agent error"))
entity = AgentEntity(mock_agent)
mock_context = Mock()
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
result = await entity.run(
mock_context, {"message": "Test message", "thread_id": "conv-1", "correlationId": "corr-app-error-1"}
)
result = await entity.run({
"message": "Test message",
"correlationId": "corr-app-error-1",
})
assert isinstance(result, AgentRunResponse)
assert len(result.messages) == 1
@@ -711,7 +724,7 @@ class TestIncomingRequestParsing:
request = Mock()
request.params = {"thread_id": "query-thread"}
req_body = {}
req_body: dict[str, Any] = {}
thread_id = app._resolve_thread_id(request, req_body)
@@ -778,7 +791,7 @@ class TestHttpRunRoute:
assert run_request["message"] == "Plain text via HTTP"
assert run_request["role"] == "user"
assert "thread_id" in run_request
assert "thread_id" not in run_request
async def test_http_run_accept_header_returns_json(self) -> None:
"""Test that Accept header requesting JSON results in JSON response."""
@@ -914,9 +927,9 @@ class TestMCPToolEndpoint:
patch.object(app, "durable_client_input") as client_mock,
):
# Setup mock decorator chain
func_name_mock.return_value = lambda f: f
mcp_trigger_mock.return_value = lambda f: f
client_mock.return_value = lambda f: f
func_name_mock.return_value = _identity_decorator
mcp_trigger_mock.return_value = _identity_decorator
client_mock.return_value = _identity_decorator
app._setup_mcp_tool_trigger(mock_agent.name, mock_agent.description)
@@ -939,11 +952,11 @@ class TestMCPToolEndpoint:
app = AgentFunctionApp()
with (
patch.object(app, "function_name", return_value=lambda f: f),
patch.object(app, "function_name", return_value=_identity_decorator),
patch.object(app, "mcp_tool_trigger") as mcp_trigger_mock,
patch.object(app, "durable_client_input", return_value=lambda f: f),
patch.object(app, "durable_client_input", return_value=_identity_decorator),
):
mcp_trigger_mock.return_value = lambda f: f
mcp_trigger_mock.return_value = _identity_decorator
app._setup_mcp_tool_trigger(mock_agent.name, None)
@@ -1065,10 +1078,10 @@ class TestMCPToolEndpoint:
app = AgentFunctionApp(agents=[mock_agent], enable_mcp_tool_trigger=True)
# Capture the health check handler function
captured_handler = None
captured_handler: Callable[[func.HttpRequest], func.HttpResponse] | None = None
def capture_decorator(*args, **kwargs):
def decorator(func):
def capture_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
def decorator(func: TFunc) -> TFunc:
nonlocal captured_handler
captured_handler = func
return func