From aa19c4f0d3284731855a2db28fd8b65866ed8020 Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Tue, 20 Oct 2026 13:59:10 +0530 Subject: [PATCH 1/2] Python: fix(core): keep AgentResponse usage when streaming a workflow agent WorkflowAgent turned an AgentResponse output into stream updates without its usage, so a streamed run lost the token counts that a non-streamed run reports. Carry the usage as usage content on the last update, the way AgentResponse.to_updates() does. When none of the response's messages is shown to the caller (only internal calls, or no messages), emit one update for the usage and the response-level fields. It has no message_id or author_name, so it joins the message in progress instead of splitting it. Fixes #9564 --- .../core/agent_framework/_workflows/_agent.py | 17 +++ .../tests/workflow/test_workflow_agent.py | 141 ++++++++++++++++++ 2 files changed, 158 insertions(+) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 4119cdafd08..2a29340f041 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -775,12 +775,29 @@ def _convert_workflow_event_to_agent_response_updates( raw_representation=msg, ) ) + if not updates and data.usage_details is not None: + # No message is shown to the caller, but the tokens were still spent. Without a + # message_id or author_name this update joins the message in progress instead of + # starting a new one. + updates.append( + AgentResponseUpdate( + contents=[], + role="assistant", + response_id=data.response_id or response_id, + created_at=data.created_at + or datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), + raw_representation=cast(AgentResponse[Any], data), + ) + ) if updates: updates[-1].agent_id = data.agent_id updates[-1].finish_reason = data.finish_reason updates[-1].continuation_token = data.continuation_token _copy_operation_state(cast(AgentResponse[Any], data), updates[-1]) updates[-1].additional_properties = dict(data.additional_properties) + # Updates carry usage as content; without it the streamed response loses the token counts. + if data.usage_details is not None: + updates[-1].contents.append(Content.from_usage(data.usage_details)) return updates if isinstance(data, Message): if data.role != "assistant" or _contains_internal_call_content(data.contents): diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 020d8a73267..4b5f82433a9 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -943,6 +943,147 @@ async def metadata_executor(messages: list[Message], ctx: WorkflowContext[Never, assert final_response.continuation_token == {"token": "response-resume-token"} assert final_response.additional_properties == {"provider_marker": "preserve-response"} + async def test_workflow_as_agent_stream_preserves_agent_response_usage(self) -> None: + """Streaming must carry the usage of an AgentResponse output, like the non-streaming path does.""" + usage = UsageDetails(input_token_count=10, output_token_count=5, total_token_count=15) + + @executor + async def usage_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output( + AgentResponse( + messages=[ + Message(role="assistant", contents=["first"]), + Message(role="assistant", contents=["second"]), + ], + response_id="usage-response", + usage_details=usage, + ) + ) + + agent = WorkflowBuilder(start_executor=usage_executor).build().as_agent("usage-test-agent") + + non_streaming = await agent.run("hello") + assert non_streaming.usage_details == usage + + stream = agent.run("hello", stream=True) + updates = [update async for update in stream] + final_response = await stream.get_final_response() + + assert [update.text for update in updates] == ["first", "second"] + # Usage describes the whole response, so it rides on the last update only. + assert not any(content.type == "usage" for content in updates[0].contents) + usage_contents = [content for content in updates[-1].contents if content.type == "usage"] + assert len(usage_contents) == 1 + assert usage_contents[0].usage_details == usage + assert final_response.usage_details == usage + assert final_response.text == non_streaming.text + + async def test_workflow_as_agent_stream_sums_usage_of_multiple_agent_responses(self) -> None: + """Usage from several AgentResponse outputs adds up the same way in both run modes.""" + + @executor + async def two_responses_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output( + AgentResponse( + messages=[Message(role="assistant", contents=["one"])], + response_id="usage-response-1", + usage_details=UsageDetails(input_token_count=3, output_token_count=2, total_token_count=5), + ) + ) + await ctx.yield_output( + AgentResponse( + messages=[Message(role="assistant", contents=["two"])], + response_id="usage-response-2", + usage_details=UsageDetails(input_token_count=7, output_token_count=1, total_token_count=8), + ) + ) + + agent = WorkflowBuilder(start_executor=two_responses_executor).build().as_agent("usage-sum-agent") + expected = UsageDetails(input_token_count=10, output_token_count=3, total_token_count=13) + + non_streaming = await agent.run("hello") + streaming = await agent.run("hello", stream=True).get_final_response() + + assert non_streaming.usage_details == expected + assert streaming.usage_details == expected + + @pytest.mark.parametrize( + "messages", + [ + [Message(role="assistant", contents=[Content.from_function_call(call_id="c1", name="t", arguments="{}")])], + [], + ], + ids=["only_internal_call", "no_messages"], + ) + async def test_workflow_as_agent_stream_keeps_usage_without_caller_facing_messages( + self, messages: list[Message] + ) -> None: + """Usage still reaches the streamed response when no message of the AgentResponse is shown to the caller.""" + usage = UsageDetails(input_token_count=10, output_token_count=5, total_token_count=15) + + @executor + async def usage_only_executor(messages_in: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output(AgentResponse(messages=list(messages), usage_details=usage)) + + agent = WorkflowBuilder(start_executor=usage_only_executor).build().as_agent("usage-only-agent") + + non_streaming = await agent.run("hello") + streaming = await agent.run("hello", stream=True).get_final_response() + + assert non_streaming.usage_details == usage + assert streaming.usage_details == usage + assert streaming.text == non_streaming.text + assert len(streaming.messages) == len(non_streaming.messages) + + async def test_workflow_as_agent_stream_usage_only_update_does_not_split_a_message(self) -> None: + """A usage-only AgentResponse between chunks of one streamed message leaves that message whole.""" + usage = UsageDetails(input_token_count=2, output_token_count=1, total_token_count=3) + + @executor + async def interleaved_executor( + messages: list[Message], + ctx: WorkflowContext[Never, AgentResponseUpdate | AgentResponse], # type: ignore[valid-type] + ) -> None: + await ctx.yield_output( + AgentResponseUpdate(contents=[Content.from_text("x")], message_id="m1", author_name="writer") + ) + await ctx.yield_output(AgentResponse(messages=[], usage_details=usage)) + await ctx.yield_output( + AgentResponseUpdate(contents=[Content.from_text("y")], message_id="m1", author_name="writer") + ) + + agent = WorkflowBuilder(start_executor=interleaved_executor).build().as_agent("interleaved-agent") + + streaming = await agent.run("hello", stream=True).get_final_response() + + assert [message.text for message in streaming.messages] == ["xy"] + assert streaming.usage_details == usage + + async def test_workflow_as_agent_stream_usage_only_update_keeps_continuation_token(self) -> None: + """The usage-only update carries the response-level fields, like the last update of a normal response.""" + usage = UsageDetails(input_token_count=4, output_token_count=0, total_token_count=4) + + @executor + async def background_executor(messages: list[Message], ctx: WorkflowContext[Never, AgentResponse]) -> None: # type: ignore[valid-type] + await ctx.yield_output( + AgentResponse( + messages=[], + usage_details=usage, + continuation_token=cast(Any, {"token": "still-running"}), + ) + ) + + agent = WorkflowBuilder(start_executor=background_executor).build().as_agent("background-agent") + + stream = agent.run("hello", stream=True) + updates = [update async for update in stream] + final_response = await stream.get_final_response() + + assert len(updates) == 1 + assert [content.type for content in updates[0].contents] == ["usage"] + assert updates[0].continuation_token == {"token": "still-running"} + assert final_response.usage_details == usage + async def test_workflow_as_agent_stream_preserves_custom_finish_reason(self) -> None: """Test that a non-literal finish_reason is forwarded unchanged. From 56d3841f444d4c0bbd1a1ee3654c32854db7a1cf Mon Sep 17 00:00:00 2001 From: Somay-kousis Date: Mon, 26 Oct 2026 21:51:38 +0530 Subject: [PATCH 2/2] Python: fix(core): don't expose filtered messages on the usage-only update The usage-only update carried the whole AgentResponse as its raw representation, so the internal messages that the caller filter drops reached the caller through it and through the final response. Leave it unset, and assert that no update or final raw representation carries the response. --- python/packages/core/agent_framework/_workflows/_agent.py | 1 - .../packages/core/tests/workflow/test_workflow_agent.py | 8 +++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 2a29340f041..e71a62a6eb3 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -786,7 +786,6 @@ def _convert_workflow_event_to_agent_response_updates( response_id=data.response_id or response_id, created_at=data.created_at or datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"), - raw_representation=cast(AgentResponse[Any], data), ) ) if updates: diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 4b5f82433a9..93a1db2a807 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -1028,12 +1028,18 @@ async def usage_only_executor(messages_in: list[Message], ctx: WorkflowContext[N agent = WorkflowBuilder(start_executor=usage_only_executor).build().as_agent("usage-only-agent") non_streaming = await agent.run("hello") - streaming = await agent.run("hello", stream=True).get_final_response() + stream = agent.run("hello", stream=True) + updates = [update async for update in stream] + streaming = await stream.get_final_response() assert non_streaming.usage_details == usage assert streaming.usage_details == usage assert streaming.text == non_streaming.text assert len(streaming.messages) == len(non_streaming.messages) + # The filtered messages must not reach the caller through a raw representation either. + assert all(update.raw_representation is None for update in updates) + raw = streaming.raw_representation + assert not any(isinstance(item, AgentResponse) for item in (raw if isinstance(raw, list) else [raw])) async def test_workflow_as_agent_stream_usage_only_update_does_not_split_a_message(self) -> None: """A usage-only AgentResponse between chunks of one streamed message leaves that message whole."""