Skip to content
Merged
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
20 changes: 12 additions & 8 deletions src/strands_compose/hooks/event_publisher.py
Original file line number Diff line number Diff line change
Expand Up @@ -252,18 +252,22 @@ def _on_complete(self, event: AfterInvocationEvent) -> None:
invocation = metrics.latest_agent_invocation
usage = invocation.usage if invocation else metrics.accumulated_usage

data: dict[str, Any] = {
"type": "agent",
"usage": {
"input_tokens": usage.get("inputTokens", 0),
"output_tokens": usage.get("outputTokens", 0),
"total_tokens": usage.get("totalTokens", 0),
},
"text": str(result) if result is not None else "",
"message": result.message if result is not None else {},
}

self._callback(
StreamEvent(
type=EventType.AGENT_COMPLETE,
agent_name=self._agent_name,
data={
"type": "agent",
"usage": {
"input_tokens": usage.get("inputTokens", 0),
"output_tokens": usage.get("outputTokens", 0),
"total_tokens": usage.get("totalTokens", 0),
},
},
data=data,
),
)

Expand Down
43 changes: 43 additions & 0 deletions tests/unit/hooks/test_event_publisher.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ def test_complete_emits_with_usage(self):
metrics.latest_agent_invocation = None
metrics.accumulated_usage = {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}
complete_event.agent.event_loop_metrics = metrics
complete_event.result = None
pub._on_complete(complete_event)

assert len(events) == 1
Expand All @@ -98,6 +99,48 @@ def test_complete_emits_with_usage(self):
assert events[0].data["usage"]["output_tokens"] == 5
assert events[0].data["usage"]["total_tokens"] == 15

def test_complete_includes_text_and_message_when_result_present(self) -> None:
"""AGENT_COMPLETE includes text and message fields when result is not None."""
events: list = []
pub = EventPublisher(callback=events.append, agent_name="test")

message = {"role": "assistant", "content": [{"text": "Hello!"}]}
result = MagicMock()
result.__str__ = MagicMock(return_value="Hello!")
result.message = message
result.stop_reason = "end_turn"

complete_event = MagicMock()
metrics = MagicMock()
metrics.latest_agent_invocation = None
metrics.accumulated_usage = {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}
complete_event.agent.event_loop_metrics = metrics
complete_event.result = result
pub._on_complete(complete_event)

assert len(events) == 1
assert events[0].type == EventType.AGENT_COMPLETE
assert events[0].data["text"] == "Hello!"
assert events[0].data["message"] == message

def test_complete_defaults_text_and_message_when_result_is_none(self) -> None:
"""AGENT_COMPLETE uses empty defaults for text and message when result is None."""
events: list = []
pub = EventPublisher(callback=events.append, agent_name="test")

complete_event = MagicMock()
metrics = MagicMock()
metrics.latest_agent_invocation = None
metrics.accumulated_usage = {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}
complete_event.agent.event_loop_metrics = metrics
complete_event.result = None
pub._on_complete(complete_event)

assert len(events) == 1
assert events[0].type == EventType.AGENT_COMPLETE
assert events[0].data["text"] == ""
assert events[0].data["message"] == {}

def test_complete_with_interrupt_result_emits_interrupt_events(self) -> None:
events = []
pub = EventPublisher(callback=events.append, agent_name="test")
Expand Down