Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions python/packages/core/agent_framework/_workflows/_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@
AgentResponseUpdate,
AgentRunInputs,
Content,
FinishReason,
FinishReasonLiteral,
Message,
ResponseStream,
UsageDetails,
Expand Down Expand Up @@ -654,9 +656,17 @@ def _convert_workflow_event_to_agent_response_updates(
contents=list(data.contents),
role=data.role,
author_name=data.author_name or executor_id,
agent_id=data.agent_id,
response_id=data.response_id,
message_id=data.message_id,
created_at=data.created_at,
# The attribute is typed wider than the constructor accepts (custom
# connectors may set any string); forward the value unchanged.
finish_reason=cast(FinishReasonLiteral | FinishReason | None, data.finish_reason),
continuation_token=data.continuation_token,
additional_properties=dict(data.additional_properties)
if data.additional_properties is not None
else None,
raw_representation=data.raw_representation,
)
]
Expand Down
105 changes: 104 additions & 1 deletion python/packages/core/tests/workflow/test_workflow_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import uuid
from collections.abc import Awaitable, Sequence
from dataclasses import dataclass
from typing import Any, Literal, overload
from typing import Any, Literal, cast, overload

import pytest
from typing_extensions import Never
Expand All @@ -15,6 +15,7 @@
AgentSession,
Content,
Executor,
FinishReason,
HistoryProvider,
InMemoryHistoryProvider,
Message,
Expand Down Expand Up @@ -852,6 +853,108 @@ async def yielding_executor(messages: list[Message], ctx: WorkflowContext[Never,
assert "first output" in texts
assert "second output" in texts

async def test_workflow_as_agent_stream_preserves_response_update_metadata(self) -> None:
"""Test that streaming forwards finish_reason, continuation_token and additional_properties.

This validates the fix for issue #7952: AgentResponseUpdate metadata should be
forwarded as-is when the workflow is wrapped via .as_agent().
"""

@executor
async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponseUpdate]) -> None: # type: ignore[valid-type]
await ctx.yield_output(
AgentResponseUpdate(
contents=[Content.from_text(text="payload")],
role="assistant",
agent_id="source-agent",
response_id="source-response",
message_id="source-message",
finish_reason="stop",
continuation_token=cast(Any, {"token": "resume-token"}),
additional_properties={"provider_marker": "preserve-me"},
)
)

workflow = WorkflowBuilder(start_executor=metadata_executor).build()
agent = workflow.as_agent("metadata-test-agent")

updates: list[AgentResponseUpdate] = []
async for update in agent.run("hello", stream=True):
updates.append(update)

metadata_updates = [u for u in updates if u.response_id == "source-response"]
assert len(metadata_updates) == 1
update = metadata_updates[0]
assert update.text == "payload"
assert update.agent_id == "source-agent"
assert update.finish_reason == "stop"
assert update.continuation_token == {"token": "resume-token"}
assert update.additional_properties == {"provider_marker": "preserve-me"}

async def test_workflow_as_agent_stream_preserves_custom_finish_reason(self) -> None:
"""Test that a non-literal finish_reason is forwarded unchanged.

Custom chat connectors can report finish reasons outside the standard
literals (modeled as `FinishReason`); the WorkflowAgent must preserve
them when re-emitting the update.
"""

@executor
async def custom_reason_executor(
messages: list[Message], ctx: WorkflowContext[Any, AgentResponseUpdate]
) -> None:
await ctx.yield_output(
AgentResponseUpdate(
contents=[Content.from_text(text="payload")],
role="assistant",
agent_id="source-agent",
response_id="custom-reason-response",
message_id="source-message",
finish_reason=FinishReason("custom_reason"),
)
)

workflow = WorkflowBuilder(start_executor=custom_reason_executor).build()
agent = workflow.as_agent("custom-reason-test-agent")

updates: list[AgentResponseUpdate] = []
async for update in agent.run("hello", stream=True):
updates.append(update)

metadata_updates = [u for u in updates if u.response_id == "custom-reason-response"]
assert len(metadata_updates) == 1
update = metadata_updates[0]
assert update.text == "payload"
assert update.agent_id == "source-agent"
assert update.finish_reason == "custom_reason"

async def test_workflow_as_agent_stream_preserves_empty_additional_properties(self) -> None:
"""Test that an explicitly empty additional_properties dict is not converted to None."""

@executor
async def empty_props_executor(
messages: list[Message], ctx: WorkflowContext[Any, AgentResponseUpdate]
) -> None:
await ctx.yield_output(
AgentResponseUpdate(
contents=[Content.from_text(text="payload")],
role="assistant",
response_id="empty-props-response",
additional_properties={},
)
)

workflow = WorkflowBuilder(start_executor=empty_props_executor).build()
agent = workflow.as_agent("empty-props-test-agent")

updates: list[AgentResponseUpdate] = []
async for update in agent.run("hello", stream=True):
updates.append(update)

forwarded = [u for u in updates if u.response_id == "empty-props-response"]
assert len(forwarded) == 1
assert forwarded[0].additional_properties == {}

async def test_workflow_as_agent_yield_output_with_content_types(self) -> None:
"""Test that yield_output preserves different content types (Content, Content, etc.)."""

Expand Down
Loading