Working: Workflow as agent with function approval

This commit is contained in:
Tao Chen
2026-06-01 22:10:28 -07:00
Unverified
parent 9c0623ca5f
commit 5fade9e2a8
8 changed files with 1080 additions and 65 deletions
@@ -2134,9 +2134,7 @@ class SkillsProvider(ContextProvider):
),
FunctionTool(
name="read_skill_resource",
description=(
"Reads a resource associated with a skill, such as references, assets, or dynamic data."
),
description=("Reads a resource associated with a skill, such as references, assets, or dynamic data."),
func=_read_resource,
input_model={
"type": "object",
@@ -2173,8 +2171,7 @@ class SkillsProvider(ContextProvider):
"type": "object",
"additionalProperties": True,
"description": (
"Named arguments as key-value pairs "
'(e.g. {"length": 24, "uppercase": true}).'
'Named arguments as key-value pairs (e.g. {"length": 24, "uppercase": true}).'
),
},
{
@@ -487,19 +487,23 @@ class WorkflowAgent(BaseAgent):
def _process_request_info_event(
self,
event: WorkflowEvent[Any],
) -> tuple[Content, Content]:
"""Convert a request_info event to FunctionCallContent and FunctionApprovalRequestContent.
) -> Content:
"""Convert a request_info event to FunctionApprovalRequestContent.
Args:
event: A WorkflowEvent with type='request_info'.
Returns:
A tuple of (FunctionCallContent, FunctionApprovalRequestContent).
"""
request_id = event.request_id
if not request_id:
raise ValueError("request_info event must have a request_id")
A FunctionApprovalRequestContent.
Note:
If the event data is already a FunctionApprovalRequestContent, it will be returned as-is.
"""
if isinstance(event.data, Content) and event.data.user_input_request:
# Return the event data as-is if it's already a properly formed FunctionApprovalRequestContent
return event.data
request_id = event.request_id
args = self.RequestInfoFunctionArgs(request_id=request_id, data=event.data).to_dict()
function_call = Content.from_function_call(
@@ -507,12 +511,10 @@ class WorkflowAgent(BaseAgent):
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=args,
)
approval_request = Content.from_function_approval_request(
return Content.from_function_approval_request(
id=request_id,
function_call=function_call,
additional_properties={"request_id": request_id},
)
return function_call, approval_request
def _convert_workflow_events_to_agent_response(
self,
@@ -531,10 +533,10 @@ class WorkflowAgent(BaseAgent):
for output_event in output_events:
if output_event.type == "request_info":
function_call, approval_request = self._process_request_info_event(output_event)
approval_request = self._process_request_info_event(output_event)
messages.append(
Message(
contents=[function_call, approval_request],
contents=[approval_request],
role="assistant",
author_name=output_event.source_executor_id,
message_id=str(uuid.uuid4()),
@@ -702,10 +704,10 @@ class WorkflowAgent(BaseAgent):
]
if event.type == "request_info":
function_call, approval_request = self._process_request_info_event(event)
approval_request = self._process_request_info_event(event)
return [
AgentResponseUpdate(
contents=[function_call, approval_request],
contents=[approval_request],
role="assistant",
author_name=self.name,
response_id=response_id,
@@ -724,27 +726,30 @@ class WorkflowAgent(BaseAgent):
for message in input_messages:
for content in message.contents:
if content.type == "function_approval_response":
request_id = content.additional_properties.get("request_id")
if not request_id:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent must have a request_id in additional_properties."
)
request_id: str = content.id # pyright: ignore[reportAssignmentType]
function_call: Content = content.function_call # type: ignore[attr-defined]
# Parse the function arguments to recover request payload
arguments_payload = content.function_call.arguments # type: ignore[attr-defined, union-attr]
if isinstance(arguments_payload, str):
try:
parsed_args = self.RequestInfoFunctionArgs.from_json(arguments_payload)
except ValueError as exc:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent arguments must decode to a mapping."
) from exc
elif isinstance(arguments_payload, dict):
parsed_args = self.RequestInfoFunctionArgs.from_dict(arguments_payload)
if function_call.name != self.REQUEST_INFO_FUNCTION_NAME:
# This response is for a raw approval request that is itself already an
# approval request.
function_responses[request_id] = content
else:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent arguments must be a mapping or JSON string."
)
function_responses[request_id] = parsed_args.data
# This response is for an approval request constructed from a request_info event.
arguments_payload = content.function_call.arguments # type: ignore[attr-defined, union-attr]
if isinstance(arguments_payload, str):
try:
parsed_args = self.RequestInfoFunctionArgs.from_json(arguments_payload)
except ValueError as exc:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent arguments must decode to a mapping."
) from exc
elif isinstance(arguments_payload, dict):
parsed_args = self.RequestInfoFunctionArgs.from_dict(arguments_payload)
else:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent arguments must be a mapping or JSON string."
)
function_responses[request_id] = parsed_args.data
elif content.type == "function_result":
response_data = content.result if hasattr(content, "result") else str(content) # type: ignore[attr-defined]
function_responses[content.call_id] = response_data # pyright: ignore[reportArgumentType]
@@ -429,15 +429,18 @@ class AgentExecutor(Executor):
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
)
await ctx.yield_output(response)
# Handle any user input requests
if response.user_input_requests:
for user_input_request in response.user_input_requests:
self._pending_agent_requests[user_input_request.id] = user_input_request # type: ignore[index]
await ctx.request_info(user_input_request, Content)
await ctx.request_info(user_input_request, Content, request_id=user_input_request.id)
return None
# Only yield output if the response is complete and not waiting for user input.
# This is to avoid emmiting two events of different types ('output' and 'request_info')
# that carry the same payload.
await ctx.yield_output(response)
return response
async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUpdate]) -> AgentResponse | None:
@@ -472,9 +475,13 @@ class AgentExecutor(Executor):
)
async for update in stream:
updates.append(update)
await ctx.yield_output(update)
if update.user_input_requests:
streamed_user_input_requests.extend(update.user_input_requests)
else:
# Only yield output events for updates that do not contain user input requests.
# This is to avoid emmiting two events of different types ('output' and 'request_info')
# that carry the same payload.
await ctx.yield_output(update)
# Prefer stream finalization when available so result hooks run
# (e.g., thread conversation updates). Fall back to reconstructing from updates
@@ -509,7 +516,7 @@ class AgentExecutor(Executor):
if user_input_requests:
for user_input_request in user_input_requests:
self._pending_agent_requests[user_input_request.id] = user_input_request # type: ignore[index]
await ctx.request_info(user_input_request, Content)
await ctx.request_info(user_input_request, Content, request_id=user_input_request.id)
return None
return response
@@ -4086,8 +4086,8 @@ class TestClassSkill:
async def test_content_is_cached(self) -> None:
skill = _MinimalClassSkill()
content1 = (await skill.get_content())
content2 = (await skill.get_content())
content1 = await skill.get_content()
content2 = await skill.get_content()
assert content1 is content2
def test_resources_are_lazy_cached(self) -> None:
@@ -5587,8 +5587,8 @@ class TestInlineSkillContentCaching:
async def test_content_cached_after_first_access(self) -> None:
"""InlineSkill.content returns the same object on subsequent accesses."""
skill = InlineSkill(frontmatter=SkillFrontmatter(name="test-skill", description="Test"), instructions="Body")
first = (await skill.get_content())
second = (await skill.get_content())
first = await skill.get_content()
second = await skill.get_content()
assert first is second # Same object (cached)
assert "<name>test-skill</name>" in first
@@ -699,3 +699,171 @@ async def test_resolve_executor_kwargs_empty_per_executor_does_not_fallback_to_g
resolved = {"exec_a": {}, GLOBAL_KWARGS_KEY: {"global_key": "global_val"}}
result = executor._resolve_executor_kwargs(resolved) # pyright: ignore[reportPrivateUsage]
assert result == {}
# region Tool approval emission
class _ApprovalEmittingAgent(BaseAgent):
"""Agent that returns a single ``function_approval_request`` Content.
Used to verify that ``AgentExecutor`` does *not* surface the approval
payload via both an ``output`` event and a ``request_info`` event in the
same superstep — only the ``request_info`` event must carry it.
"""
def __init__(
self,
*,
approval_request_id: str = "apr_1",
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
**kwargs: Any,
):
super().__init__(**kwargs)
self._approval_request_id = approval_request_id
self._tool_name = tool_name
self._tool_arguments: dict[str, Any] = tool_arguments or {"path": "/tmp/secret.txt"}
self.run_count = 0
def _build_approval_content(self) -> Content:
function_call = Content.from_function_call(
call_id=self._approval_request_id,
name=self._tool_name,
arguments=self._tool_arguments,
)
return Content.from_function_approval_request(id=self._approval_request_id, function_call=function_call)
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.run_count += 1
approval = self._build_approval_content()
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(contents=[approval], role="assistant")
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
async def _run() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", [approval])])
return _run()
def _has_approval_payload(event: WorkflowEvent[Any]) -> bool:
"""Return True if the event's data carries a ``function_approval_request`` content."""
data: Any = event.data
def _contents_of(value: Any) -> list[Content]:
if isinstance(value, AgentResponseUpdate):
return list(value.contents)
if isinstance(value, AgentResponse):
return [c for m in value.messages for c in m.contents]
if isinstance(value, AgentExecutorResponse):
return [c for m in value.agent_response.messages for c in m.contents]
if isinstance(value, Message):
return list(value.contents)
if isinstance(value, Content):
return [value]
return []
return any(c.type == "function_approval_request" for c in _contents_of(data))
async def test_agent_executor_does_not_double_emit_approval_non_streaming() -> None:
"""Non-streaming: approval payload must only appear in the ``request_info`` event.
Regression test for the bug where ``AgentExecutor._run_agent`` first
``yield_output``-ed the response (carrying the approval Content) and then
additionally emitted a ``request_info`` event for the same payload.
"""
agent = _ApprovalEmittingAgent(id="approve_agent", name="ApproveAgent", approval_request_id="apr_ns_1")
executor = AgentExecutor(agent, id="approve_exec")
workflow = WorkflowBuilder(start_executor=executor).build()
request_info_events: list[WorkflowEvent[Any]] = []
output_events: list[WorkflowEvent[Any]] = []
for event in await workflow.run("please delete it"):
if event.type == "request_info":
request_info_events.append(event)
elif event.type == "output":
output_events.append(event)
assert len(request_info_events) == 1
assert _has_approval_payload(request_info_events[0])
# The approval payload must not also be surfaced as a workflow output.
assert not any(_has_approval_payload(e) for e in output_events)
assert agent.run_count == 1
async def test_agent_executor_does_not_double_emit_approval_streaming() -> None:
"""Streaming: per-update approval payload must not be ``yield_output``-ed."""
agent = _ApprovalEmittingAgent(id="approve_agent_s", name="ApproveAgentS", approval_request_id="apr_st_1")
executor = AgentExecutor(agent, id="approve_exec_s")
workflow = WorkflowBuilder(start_executor=executor).build()
request_info_events: list[WorkflowEvent[Any]] = []
output_events: list[WorkflowEvent[Any]] = []
async for event in workflow.run("please delete it", stream=True):
if event.type == "request_info":
request_info_events.append(event)
elif event.type == "output":
output_events.append(event)
assert len(request_info_events) == 1
assert _has_approval_payload(request_info_events[0])
assert not any(_has_approval_payload(e) for e in output_events)
assert agent.run_count == 1
async def test_agent_executor_request_info_uses_user_input_request_id() -> None:
"""``ctx.request_info`` must register the request under the agent's approval id.
This makes the workflow's pending-request id round-trip with the
``function_approval_response.id`` the caller echoes back, so
``Workflow._send_responses_internal`` can look it up directly.
"""
agent = _ApprovalEmittingAgent(id="approve_agent_id", name="ApproveAgentId", approval_request_id="apr_match")
executor = AgentExecutor(agent, id="approve_exec_id")
workflow = WorkflowBuilder(start_executor=executor).build()
request_info_events: list[WorkflowEvent[Any]] = []
async for event in workflow.run("please delete it", stream=True):
if event.type == "request_info":
request_info_events.append(event)
assert len(request_info_events) == 1
assert request_info_events[0].request_id == "apr_match"
# endregion Tool approval emission
@@ -245,10 +245,13 @@ class TestWorkflowAgent:
assert approval_update is not None, "Should have received a request_info approval request"
function_call = next(content for content in approval_update.contents if content.type == "function_call")
approval_request = next(
content for content in approval_update.contents if content.type == "function_approval_request"
)
assert approval_request.id is not None
assert approval_request.function_call is not None
function_call = approval_request.function_call
# Verify the function call has expected structure
assert function_call.call_id is not None
@@ -256,12 +259,6 @@ class TestWorkflowAgent:
assert isinstance(function_call.arguments, dict)
assert function_call.arguments.get("request_id") == approval_request.id
# Approval request should reference the same function call
assert approval_request.id is not None
assert approval_request.function_call is not None
assert approval_request.function_call.call_id == function_call.call_id
assert approval_request.function_call.name == function_call.name
# Verify the request is tracked in pending_requests
pending_requests = await workflow._runner_context.get_pending_request_info_events()
assert len(pending_requests) == 1
@@ -1564,3 +1561,350 @@ class TestWorkflowAgentMergeUpdates:
# Order: text (user), text (assistant), function_result (orphan at end)
assert content_types == ["text", "text", "function_result"]
class _ToolApprovalMockAgent(SupportsAgentRun):
"""Mock agent whose first run returns a FunctionApprovalRequestContent.
Subsequent runs (after receiving an approval response in the input messages)
return a final assistant text response that echoes the approved arguments.
This mirrors a real agent whose tool invocation requires user approval.
"""
def __init__(
self,
name: str,
*,
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
approval_request_ids: Sequence[str] | None = None,
) -> None:
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._tool_name = tool_name
self._tool_arguments = tool_arguments or {"path": "/tmp/example"}
# Pre-allocated request ids so the test can verify what the WorkflowAgent forwards.
self._approval_request_ids: list[str] = list(approval_request_ids) if approval_request_ids else []
self.run_count = 0
# Inputs received on the most recent (continuation) run, for assertions.
self.last_run_messages: list[Message] = []
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
return AgentSession()
def _next_request_id(self) -> str:
if self._approval_request_ids:
return self._approval_request_ids.pop(0)
return str(uuid.uuid4())
def _build_approval_request(self) -> Content:
request_id = self._next_request_id()
function_call = Content.from_function_call(
call_id=request_id,
name=self._tool_name,
arguments=self._tool_arguments,
)
return Content.from_function_approval_request(id=request_id, function_call=function_call)
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
return self._run_stream(messages=messages, session=session, **kwargs)
return self._run(messages=messages, session=session, **kwargs)
def _normalize(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None,
) -> list[Message]:
if messages is None:
return []
if isinstance(messages, str):
return [Message(role="user", contents=[Content.from_text(text=messages)])]
if isinstance(messages, Message):
return [messages]
if isinstance(messages, Content):
return [Message(role="user", contents=[messages])]
result: list[Message] = []
for item in messages:
if isinstance(item, Message):
result.append(item)
elif isinstance(item, Content):
result.append(Message(role="user", contents=[item]))
else:
result.append(Message(role="user", contents=[Content.from_text(text=item)]))
return result
def _approval_responses_in(self, messages: list[Message]) -> list[Content]:
approvals: list[Content] = []
for msg in messages:
for content in msg.contents:
if content.type == "function_approval_response":
approvals.append(content)
return approvals
async def _run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
approvals = self._approval_responses_in(normalized)
if approvals:
# Continuation: reflect approved arguments in the final response text.
approved_text = "; ".join(
f"approved={a.approved} id={a.id}" # type: ignore[attr-defined]
for a in approvals
)
return AgentResponse(messages=[Message("assistant", [Content.from_text(text=f"done ({approved_text})")])])
# First run: ask for tool approval.
approval = self._build_approval_request()
return AgentResponse(messages=[Message("assistant", [approval])])
def _run_stream(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
approvals = self._approval_responses_in(normalized)
async def _iter():
if approvals:
approved_text = "; ".join(
f"approved={a.approved} id={a.id}" # type: ignore[attr-defined]
for a in approvals
)
yield AgentResponseUpdate(
contents=[Content.from_text(text=f"done ({approved_text})")],
role="assistant",
author_name=self.name,
)
return
approval = self._build_approval_request()
yield AgentResponseUpdate(
contents=[approval],
role="assistant",
author_name=self.name,
)
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
class TestWorkflowAgentToolApproval:
"""Tests for tool-approval requests bubbling through WorkflowAgent.
Covers the case where a workflow contains an AgentExecutor whose underlying
agent emits a FunctionApprovalRequestContent (tool needing user approval).
The WorkflowAgent must:
* forward the original FunctionApprovalRequestContent unchanged (no
wrapping inside a synthesized 'request_info' function call), and
* route a subsequent FunctionApprovalResponseContent back to the
AgentExecutor so the agent can resume.
"""
def _find_approval_request(
self,
contents: Sequence[Content],
tool_name: str,
) -> Content | None:
for content in contents:
if (
content.type == "function_approval_request"
and getattr(content.function_call, "name", None) == tool_name # type: ignore[attr-defined]
):
return content
return None
async def test_tool_approval_request_forwarded_unchanged(self) -> None:
"""The agent's FunctionApprovalRequestContent surfaces verbatim (not re-wrapped)."""
approval_id = "approval-abc-123"
mock_agent = _ToolApprovalMockAgent(
name="approval-agent",
tool_name="delete_file",
tool_arguments={"path": "/tmp/secret.txt"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Test Agent")
result = await agent.run("please delete the file")
assert isinstance(result, AgentResponse)
# Locate the approval request emitted by the WorkflowAgent.
all_contents: list[Content] = [c for m in result.messages for c in m.contents]
approval = self._find_approval_request(all_contents, tool_name="delete_file")
assert approval is not None, "WorkflowAgent did not forward the tool approval request"
# The id and inner function_call must match what the underlying agent produced
# — i.e. the WorkflowAgent must NOT have re-wrapped it inside a synthesized
# 'request_info' approval request.
assert approval.id == approval_id
function_call = approval.function_call # type: ignore[attr-defined]
assert function_call is not None
assert function_call.name == "delete_file"
assert function_call.name != WorkflowAgent.REQUEST_INFO_FUNCTION_NAME
assert function_call.arguments == {"path": "/tmp/secret.txt"}
# The agent must be paused awaiting the approval response.
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id in pending
async def test_tool_approval_request_forwarded_unchanged_streaming(self) -> None:
"""Streaming variant: the approval request is forwarded as-is in updates."""
approval_id = "approval-stream-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-agent-stream",
tool_name="send_email",
tool_arguments={"to": "alice@example.com"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Stream Agent")
updates: list[AgentResponseUpdate] = []
async for update in agent.run("hi", stream=True):
updates.append(update)
approval_updates = [u for u in updates if any(c.type == "function_approval_request" for c in u.contents)]
assert approval_updates, "Streaming did not surface a tool approval request"
approval = self._find_approval_request(approval_updates[-1].contents, tool_name="send_email")
assert approval is not None
assert approval.id == approval_id
function_call = approval.function_call # type: ignore[attr-defined]
assert function_call is not None
assert function_call.name == "send_email"
assert function_call.name != WorkflowAgent.REQUEST_INFO_FUNCTION_NAME
assert function_call.arguments == {"to": "alice@example.com"}
async def test_tool_approval_response_resumes_agent(self) -> None:
"""Sending the approval response back resumes the agent and clears pending requests."""
approval_id = "approval-resume-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-resume-agent",
tool_name="delete_file",
tool_arguments={"path": "/tmp/x"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Resume Agent")
first_result = await agent.run("delete it")
approval = self._find_approval_request(
[c for m in first_result.messages for c in m.contents],
tool_name="delete_file",
)
assert approval is not None
assert mock_agent.run_count == 1
# Build the approval response. NOTE: the inner function_call's name is the
# original tool name ('delete_file'), NOT 'request_info'. This exercises the
# branch in WorkflowAgent._extract_function_responses that routes raw
# tool-approval responses straight through using content.id.
approval_response = approval.to_function_approval_response(approved=True) # type: ignore[attr-defined]
response_message = Message(role="user", contents=[approval_response])
final_result = await agent.run(response_message)
assert isinstance(final_result, AgentResponse)
# The mock agent should have been invoked a second time and seen the
# approval response in its inputs.
assert mock_agent.run_count == 2
approvals_seen = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approvals_seen) == 1
assert approvals_seen[0].id == approval_id # type: ignore[attr-defined]
assert approvals_seen[0].approved is True # type: ignore[attr-defined]
# The pending approval should now be cleared.
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id not in pending
# The final assistant message reflects the resumption.
final_text = " ".join(m.text or "" for m in final_result.messages)
assert "done" in final_text
assert approval_id in final_text
async def test_tool_approval_request_id_matches_pending_request(self) -> None:
"""The approval request id surfaced by WorkflowAgent matches the workflow's pending request id.
This guards the AgentExecutor change that forwards
request_id=user_input_request.id to ctx.request_info(...), which is what
allows the response routed back via WorkflowAgent to resolve the pending
request without an id-mismatch error.
"""
approval_id = "approval-id-match-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-id-match-agent",
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Id Agent")
await agent.run("go")
pending = await workflow._runner_context.get_pending_request_info_events()
# The agent's approval id is used as the workflow's pending request id.
assert list(pending.keys()) == [approval_id]
@@ -11,24 +11,33 @@ the registered _handle_create handler.
from __future__ import annotations
import json
from collections.abc import AsyncIterator, Callable
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import Literal, overload
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from agent_framework import (
AgentExecutorRequest,
AgentResponse,
AgentResponseUpdate,
AgentSession,
Content,
FileCheckpointStorage,
HistoryProvider,
Message,
RawAgent,
ResponseStream,
SupportsAgentRun,
WorkflowAgent,
WorkflowBuilder,
WorkflowCheckpoint,
WorkflowCheckpointException,
WorkflowContext,
WorkflowMessage,
executor,
)
from azure.ai.agentserver.responses import InMemoryResponseProvider
from mcp import McpError
@@ -101,7 +110,7 @@ def _make_agent(
return agent
def _make_server(agent: MagicMock, **kwargs: Any) -> ResponsesHostServer:
def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer:
"""Create a ResponsesHostServer with an in-memory store."""
return ResponsesHostServer(agent, store=InMemoryResponseProvider(), **kwargs)
@@ -3248,3 +3257,497 @@ class TestOAuthConsentSurfacing:
# endregion
# region Workflow agent hosting (end-to-end)
class _ToolApprovalWorkflowAgentMock(SupportsAgentRun):
"""Inner agent for a hosted ``WorkflowAgent`` whose first run emits a
``FunctionApprovalRequestContent`` and whose follow-up run (after
receiving a ``FunctionApprovalResponseContent`` in its inputs) returns a
final assistant text response.
Mirrors a real agent whose tool invocation requires user approval. Used
here to exercise the full HTTP pipeline through ``ResponsesHostServer``
when the hosted agent is a ``WorkflowAgent`` containing a tool-approval
flow.
"""
def __init__(
self,
name: str,
*,
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
approval_request_ids: Sequence[str] | None = None,
final_text: str = "done",
) -> None:
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._tool_name = tool_name
self._tool_arguments = tool_arguments or {"path": "/tmp/example"}
self._approval_request_ids: list[str] = list(approval_request_ids) if approval_request_ids else []
self._final_text = final_text
self.run_count = 0
self.last_run_messages: list[Message] = []
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
return AgentSession()
def _next_request_id(self) -> str:
# Stable across calls: when the workflow checkpoint round-trips through
# restore, ``AgentExecutor`` re-invokes the inner agent during replay.
# We must surface the *same* approval request id on each invocation so
# the workflow's pending-request id matches the id the test echoes
# back as ``mcp_approval_response``.
if self._approval_request_ids:
return self._approval_request_ids[0]
return str(uuid.uuid4())
def _build_approval_request(self) -> Content:
request_id = self._next_request_id()
function_call = Content.from_function_call(
call_id=request_id,
name=self._tool_name,
arguments=self._tool_arguments,
additional_properties={"server_label": "test_server"},
)
return Content.from_function_approval_request(id=request_id, function_call=function_call)
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
return self._run_stream(messages=messages, **kwargs)
return self._run(messages=messages, **kwargs)
@staticmethod
def _normalize(
messages: str | Content | Message | Sequence[str | Content | Message] | None,
) -> list[Message]:
if messages is None:
return []
if isinstance(messages, str):
return [Message(role="user", contents=[Content.from_text(text=messages)])]
if isinstance(messages, Message):
return [messages]
if isinstance(messages, Content):
return [Message(role="user", contents=[messages])]
result: list[Message] = []
for item in messages:
if isinstance(item, Message):
result.append(item)
elif isinstance(item, Content):
result.append(Message(role="user", contents=[item]))
else:
result.append(Message(role="user", contents=[Content.from_text(text=item)]))
return result
@staticmethod
def _approval_responses_in(messages: list[Message]) -> list[Content]:
return [c for m in messages for c in m.contents if c.type == "function_approval_response"]
async def _run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
**kwargs: Any,
) -> AgentResponse:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
if self._approval_responses_in(normalized):
return AgentResponse(messages=[Message("assistant", [Content.from_text(text=self._final_text)])])
approval = self._build_approval_request()
return AgentResponse(messages=[Message("assistant", [approval])])
def _run_stream(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
approvals = self._approval_responses_in(normalized)
async def _iter() -> AsyncIterator[AgentResponseUpdate]:
if approvals:
yield AgentResponseUpdate(
contents=[Content.from_text(text=self._final_text)],
role="assistant",
author_name=self.name,
)
return
yield AgentResponseUpdate(
contents=[self._build_approval_request()],
role="assistant",
author_name=self.name,
)
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
def _build_text_workflow_agent(text: str) -> WorkflowAgent:
"""Build a minimal ``WorkflowAgent`` whose inner agent emits a fixed text."""
class _TextAgent(SupportsAgentRun):
def __init__(self, name: str, text: str) -> None:
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._text = text
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
return AgentSession()
@overload
def run(
self,
messages: Any = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: Any = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: Any = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
text = self._text
name = self.name
async def _aresult() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", [Content.from_text(text=text)])])
async def _aiter() -> AsyncIterator[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[Content.from_text(text=text)],
role="assistant",
author_name=name,
)
if stream:
return ResponseStream(_aiter(), finalizer=AgentResponse.from_updates)
return _aresult()
inner = _TextAgent("text-agent", text)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, inner).build()
return WorkflowAgent(workflow=workflow, name="Text Workflow Agent")
def _build_approval_workflow_agent(
*,
approval_request_id: str,
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
final_text: str = "done",
) -> tuple[WorkflowAgent, _ToolApprovalWorkflowAgentMock]:
"""Build a ``WorkflowAgent`` whose inner agent emits a tool approval request."""
mock_agent = _ToolApprovalWorkflowAgentMock(
name="approval-agent",
tool_name=tool_name,
tool_arguments=tool_arguments or {"path": "/tmp/secret.txt"},
approval_request_ids=[approval_request_id],
final_text=final_text,
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
workflow_agent = WorkflowAgent(workflow=workflow, name="Approval Workflow Agent")
return workflow_agent, mock_agent
class TestWorkflowAgentHosting:
"""End-to-end HTTP tests for ``ResponsesHostServer`` hosting a ``WorkflowAgent``.
These tests drive ``_handle_inner_workflow`` through the ASGI stack:
they exercise checkpoint write/restore (multi-turn) and the
tool-approval round-trip path, which is the primary differentiator
relative to the regular agent path.
"""
async def test_basic_text_response(self) -> None:
workflow_agent = _build_text_workflow_agent("hello from workflow")
server = _make_server(workflow_agent)
resp = await _post(server, input_text="hi", stream=False)
assert resp.status_code == 200
body = resp.json()
assert body["status"] == "completed"
text_found = any(
part.get("type") == "output_text" and part.get("text") == "hello from workflow"
for item in body["output"]
if item["type"] == "message"
for part in item.get("content", [])
)
assert text_found, f"Expected workflow output text in {body['output']}"
async def test_basic_text_response_streaming(self) -> None:
workflow_agent = _build_text_workflow_agent("hello stream")
server = _make_server(workflow_agent)
resp = await _post(server, input_text="hi", stream=True)
assert resp.status_code == 200
events = _parse_sse_events(resp.text)
types = _sse_event_types(events)
assert types[0] == "response.created"
assert types[-1] == "response.completed"
assert "response.output_text.delta" in types
text_done = [e for e in events if e["event"] == "response.output_text.done"]
assert any(e["data"]["text"] == "hello stream" for e in text_done)
async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None:
workflow_agent, mock_agent = _build_approval_workflow_agent(approval_request_id="apr_wf_ns")
server = _make_server(workflow_agent)
resp = await _post(server, stream=False)
assert resp.status_code == 200
body = resp.json()
assert body["status"] == "completed"
approval_items = [it for it in body["output"] if it["type"] == "mcp_approval_request"]
assert len(approval_items) == 1
assert approval_items[0]["name"] == "delete_file"
assert approval_items[0]["server_label"] == "test_server"
approval_request_id = approval_items[0]["id"]
# The id surfaced over the wire is generated by the response stream
# builder; the original approval ``Content`` (carrying the inner
# ``function_call``) must be persisted under that id so the next
# turn can reconstruct it.
loaded = await server._approval_storage.load_approval_request( # pyright: ignore[reportPrivateUsage]
approval_request_id
)
assert loaded.type == "function_approval_request"
assert loaded.function_call.name == "delete_file" # type: ignore[attr-defined]
assert mock_agent.run_count == 1
async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None:
workflow_agent, mock_agent = _build_approval_workflow_agent(approval_request_id="apr_wf_st")
server = _make_server(workflow_agent)
resp = await _post(server, stream=True)
assert resp.status_code == 200
events = _parse_sse_events(resp.text)
types = _sse_event_types(events)
assert types[0] == "response.created"
assert types[-1] == "response.completed"
approval_request_id: str | None = None
for e in events:
if e["event"] != "response.output_item.added":
continue
item = e["data"].get("item") or {}
if item.get("type") == "mcp_approval_request":
approval_request_id = item.get("id")
break
assert approval_request_id is not None
loaded = await server._approval_storage.load_approval_request( # pyright: ignore[reportPrivateUsage]
approval_request_id
)
assert loaded.type == "function_approval_request"
assert mock_agent.run_count == 1
async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None:
"""Two-turn HTTP round-trip:
Turn 1 emits ``mcp_approval_request`` and writes a workflow
checkpoint under the response id. Turn 2 sends the
``mcp_approval_response`` with ``previous_response_id`` set, so the
host restores the checkpoint, the WorkflowAgent routes the
approval response back to the paused inner agent, and the inner
agent emits the final assistant text.
"""
workflow_agent, mock_agent = _build_approval_workflow_agent(
approval_request_id="apr_wf_rt",
final_text="done with approval",
)
server = _make_server(workflow_agent)
first = await _post(server, stream=False)
assert first.status_code == 200
first_body = first.json()
first_response_id = first_body["id"]
approval_items = [it for it in first_body["output"] if it["type"] == "mcp_approval_request"]
assert len(approval_items) == 1
approval_request_id = approval_items[0]["id"]
assert mock_agent.run_count == 1
second_payload: dict[str, Any] = {
"model": "test-model",
"input": [
{
"type": "mcp_approval_response",
"approval_request_id": approval_request_id,
"approve": True,
}
],
"stream": False,
"previous_response_id": first_response_id,
}
second = await _post_json(server, second_payload)
assert second.status_code == 200
second_body = second.json()
assert second_body["status"] == "completed"
# The inner agent must have been resumed (restore replay + new turn).
# Restore call is a no-op for the mock (no input); the new-turn call
# delivers the approval response, so run_count grows by at least 1.
assert mock_agent.run_count >= 2
# The final assistant text from the resumed inner agent surfaces in
# the HTTP output.
text_pieces = [
part.get("text", "")
for item in second_body["output"]
if item["type"] == "message"
for part in item.get("content", [])
if part.get("type") == "output_text"
]
assert any("done with approval" in t for t in text_pieces), (
f"expected resumed workflow output, got {second_body['output']}"
)
# The new-turn invocation of the inner agent must have received the
# approval response routed back through WorkflowAgent.
approval_responses = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approval_responses) == 1
assert approval_responses[0].approved is True # type: ignore[attr-defined]
async def test_round_trip_approval_response_streaming(self) -> None:
"""Streaming variant of the round-trip: turn 2 is requested with
``stream=true`` and surfaces the resumed text as SSE events."""
workflow_agent, mock_agent = _build_approval_workflow_agent(
approval_request_id="apr_wf_rt_st",
final_text="streamed-done",
)
server = _make_server(workflow_agent)
first = await _post(server, stream=False)
first_body = first.json()
first_response_id = first_body["id"]
approval_request_id = next(it["id"] for it in first_body["output"] if it["type"] == "mcp_approval_request")
second = await _post_json(
server,
{
"model": "test-model",
"input": [
{
"type": "mcp_approval_response",
"approval_request_id": approval_request_id,
"approve": True,
}
],
"stream": True,
"previous_response_id": first_response_id,
},
)
assert second.status_code == 200
events = _parse_sse_events(second.text)
types = _sse_event_types(events)
assert types[0] == "response.created"
assert types[-1] == "response.completed"
text_done = [e for e in events if e["event"] == "response.output_text.done"]
assert any("streamed-done" in e["data"]["text"] for e in text_done)
assert mock_agent.run_count >= 2
async def test_round_trip_approval_response_rejected(self) -> None:
"""Sending ``approve=False`` must surface as ``approved=False`` to the
inner agent on resume."""
workflow_agent, mock_agent = _build_approval_workflow_agent(
approval_request_id="apr_wf_reject",
final_text="acknowledged",
)
server = _make_server(workflow_agent)
first = await _post(server, stream=False)
first_body = first.json()
first_response_id = first_body["id"]
approval_request_id = next(it["id"] for it in first_body["output"] if it["type"] == "mcp_approval_request")
second = await _post_json(
server,
{
"model": "test-model",
"input": [
{
"type": "mcp_approval_response",
"approval_request_id": approval_request_id,
"approve": False,
}
],
"stream": False,
"previous_response_id": first_response_id,
},
)
assert second.status_code == 200
approval_responses = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approval_responses) == 1
assert approval_responses[0].approved is False # type: ignore[attr-defined]
# endregion
@@ -1,9 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.
import os
from random import randint
from agent_framework import Agent, AgentExecutor, WorkflowBuilder, tool
from agent_framework import Agent, AgentExecutor, WorkflowBuilder
from agent_framework.foundry import FoundryChatClient
from agent_framework_foundry_hosting import ResponsesHostServer
from azure.identity import DefaultAzureCredential
@@ -13,13 +12,6 @@ from dotenv import load_dotenv
load_dotenv()
@tool(approval_mode="always_require")
def get_locale() -> str:
"""Get the weather for a given location."""
conditions = ["CZ", "US", "FR", "DE"]
return f"The locale is {conditions[randint(0, 3)]}."
def main():
client = FoundryChatClient(
project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"],
@@ -30,7 +22,6 @@ def main():
writer_agent = Agent(
client=client,
instructions=("You are an excellent slogan writer. You create new slogans based on the given topic."),
tools=[get_locale],
name="writer",
)