Python: Fix workflow samples for bugbash: part 1 (#4055)

* Fix workflow samples for bugbash: part 1

* Fix mypy

* Fix tests
This commit is contained in:
Tao Chen
2026-02-18 15:08:53 -08:00
committed by GitHub
Unverified
parent 2dfe90306b
commit 7cee839982
11 changed files with 132 additions and 113 deletions
@@ -30,6 +30,7 @@ Key properties:
"""
import inspect
import json
import logging
import sys
from collections.abc import Awaitable, Callable, Sequence
@@ -139,7 +140,10 @@ class _AutoHandoffMiddleware(FunctionMiddleware):
from agent_framework._middleware import MiddlewareTermination
# Short-circuit execution and provide deterministic response payload for the tool call.
context.result = {HANDOFF_FUNCTION_RESULT_KEY: self._handoff_functions[context.function.name]}
# Parse the result using the default parser to ensure in a form that can be passed directly to LLM APIs.
context.result = FunctionTool.parse_result({
HANDOFF_FUNCTION_RESULT_KEY: self._handoff_functions[context.function.name]
})
raise MiddlewareTermination(result=context.result)
@@ -493,9 +497,22 @@ class HandoffAgentExecutor(AgentExecutor):
last_message = response.messages[-1]
for content in last_message.contents:
if content.type == "function_result":
# Use string comparison instead of isinstance to improve performance
if content.result and isinstance(content.result, dict):
handoff_target = content.result.get(HANDOFF_FUNCTION_RESULT_KEY) # type: ignore
if not content.result:
continue
parsed_result: dict[str, Any] | None = None
if isinstance(content.result, dict):
parsed_result = content.result
elif isinstance(content.result, str):
try:
loaded_result = json.loads(content.result)
except json.JSONDecodeError:
continue
if isinstance(loaded_result, dict):
parsed_result = loaded_result
if parsed_result is not None:
handoff_target = parsed_result.get(HANDOFF_FUNCTION_RESULT_KEY)
if isinstance(handoff_target, str):
return handoff_target
else:
@@ -17,10 +17,17 @@ from agent_framework import (
resolve_agent_id,
)
from agent_framework._clients import BaseChatClient
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._tools import FunctionInvocationLayer
from agent_framework._middleware import ChatMiddlewareLayer, FunctionInvocationContext, MiddlewareTermination
from agent_framework._tools import FunctionInvocationLayer, FunctionTool, tool
from agent_framework.orchestrations import HandoffAgentUserRequest, HandoffBuilder
from agent_framework_orchestrations._handoff import (
HANDOFF_FUNCTION_RESULT_KEY,
HandoffConfiguration,
_AutoHandoffMiddleware, # pyright: ignore[reportPrivateUsage]
get_handoff_tool_name,
)
class MockChatClient(ChatMiddlewareLayer[Any], FunctionInvocationLayer[Any], BaseChatClient[Any]):
"""Mock chat client for testing handoff workflows."""
@@ -365,3 +372,41 @@ def test_handoff_builder_accepts_all_instances_in_add_handoff():
assert "triage" in workflow.executors
assert "specialist_a" in workflow.executors
assert "specialist_b" in workflow.executors
async def test_auto_handoff_middleware_intercepts_handoff_tool_call() -> None:
"""Middleware should short-circuit matching handoff tool calls with a synthetic result."""
target_id = "specialist"
middleware = _AutoHandoffMiddleware([HandoffConfiguration(target=target_id)])
@tool(name=get_handoff_tool_name(target_id), approval_mode="never_require")
def handoff_tool() -> str:
return "unreachable"
context = FunctionInvocationContext(function=handoff_tool, arguments={})
call_next = AsyncMock()
with pytest.raises(MiddlewareTermination) as exc_info:
await middleware.process(context, call_next)
call_next.assert_not_awaited()
expected_result = FunctionTool.parse_result({HANDOFF_FUNCTION_RESULT_KEY: target_id})
assert context.result == expected_result
assert exc_info.value.result == expected_result
async def test_auto_handoff_middleware_calls_next_for_non_handoff_tool() -> None:
"""Middleware should pass through when the function name is not a configured handoff tool."""
middleware = _AutoHandoffMiddleware([HandoffConfiguration(target="specialist")])
@tool(name="regular_tool", approval_mode="never_require")
def regular_tool() -> str:
return "ok"
context = FunctionInvocationContext(function=regular_tool, arguments={})
call_next = AsyncMock()
await middleware.process(context, call_next)
call_next.assert_awaited_once()
assert context.result is None