diff --git a/src/strands/agent/agent.py b/src/strands/agent/agent.py index 7c63c1e89..853aa3108 100644 --- a/src/strands/agent/agent.py +++ b/src/strands/agent/agent.py @@ -731,6 +731,7 @@ async def _run_loop( """ self.hooks.invoke_callbacks(BeforeInvocationEvent(agent=self)) + agent_result: Optional[AgentResult] = None try: yield InitEventLoopEvent() @@ -759,9 +760,13 @@ async def _run_loop( self._session_manager.redact_latest_message(self.messages[-1], self) yield event + # Capture the result from the final event if available + if hasattr(event, "__getitem__") and "stop" in event: + agent_result = AgentResult(*event["stop"]) + finally: self.conversation_manager.apply_management(self) - self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self)) + self.hooks.invoke_callbacks(AfterInvocationEvent(agent=self, result=agent_result)) async def _execute_event_loop_cycle( self, invocation_state: dict[str, Any], structured_output_context: StructuredOutputContext | None = None diff --git a/src/strands/hooks/events.py b/src/strands/hooks/events.py index 05be255f6..778af62c2 100644 --- a/src/strands/hooks/events.py +++ b/src/strands/hooks/events.py @@ -5,10 +5,13 @@ import uuid from dataclasses import dataclass -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from typing_extensions import override +if TYPE_CHECKING: + from ..agent.agent_result import AgentResult + from ..types.content import Message from ..types.interrupt import _Interruptible from ..types.streaming import StopReason @@ -60,8 +63,15 @@ class AfterInvocationEvent(HookEvent): - Agent.__call__ - Agent.stream_async - Agent.structured_output + + Attributes: + result: The result of the agent invocation, if available. This will be None + when invoked from structured_output methods, as those return typed output + directly rather than AgentResult. """ + result: Optional["AgentResult"] = None + @property def should_reverse_callbacks(self) -> bool: """True to invoke callbacks in reverse order.""" diff --git a/tests/strands/agent/hooks/test_events.py b/tests/strands/agent/hooks/test_events.py index 8bbd89c17..d7b2f2b30 100644 --- a/tests/strands/agent/hooks/test_events.py +++ b/tests/strands/agent/hooks/test_events.py @@ -2,6 +2,7 @@ import pytest +from strands.agent.agent_result import AgentResult from strands.hooks import ( AfterInvocationEvent, AfterToolCallEvent, @@ -10,6 +11,7 @@ BeforeToolCallEvent, MessageAddedEvent, ) +from strands.types.content import Message from strands.types.tools import ToolResult, ToolUse @@ -138,3 +140,46 @@ def test_after_tool_invocation_event_cannot_write_properties(after_tool_event): after_tool_event.invocation_state = {} with pytest.raises(AttributeError, match="Property exception is not writable"): after_tool_event.exception = Exception("test") + + +def test_after_invocation_event_has_optional_result(agent): + """Test that AfterInvocationEvent has optional result field.""" + # Test with no result (structured_output case) + event_without_result = AfterInvocationEvent(agent=agent) + assert event_without_result.result is None + + # Test with result (normal invocation case) + mock_message: Message = {"role": "assistant", "content": [{"text": "test"}]} + mock_result = AgentResult( + stop_reason="end_turn", + message=mock_message, + metrics={}, + state={}, + ) + event_with_result = AfterInvocationEvent(agent=agent, result=mock_result) + assert event_with_result.result == mock_result + assert event_with_result.result.stop_reason == "end_turn" + + +def test_after_invocation_event_result_not_writable(agent): + """Test that result property is not writable after initialization.""" + mock_message: Message = {"role": "assistant", "content": [{"text": "test"}]} + mock_result = AgentResult( + stop_reason="end_turn", + message=mock_message, + metrics={}, + state={}, + ) + + event = AfterInvocationEvent(agent=agent, result=None) + + with pytest.raises(AttributeError, match="Property result is not writable"): + event.result = mock_result + + +def test_after_invocation_event_agent_not_writable(agent): + """Test that agent property is not writable.""" + event = AfterInvocationEvent(agent=agent) + + with pytest.raises(AttributeError, match="Property agent is not writable"): + event.agent = Mock() diff --git a/tests/strands/agent/test_agent_hooks.py b/tests/strands/agent/test_agent_hooks.py index 32266c3eb..2a3728070 100644 --- a/tests/strands/agent/test_agent_hooks.py +++ b/tests/strands/agent/test_agent_hooks.py @@ -197,7 +197,10 @@ def test_agent__call__hooks(agent, hook_provider, agent_tool, mock_model, tool_u ) assert next(events) == MessageAddedEvent(agent=agent, message=agent.messages[3]) - assert next(events) == AfterInvocationEvent(agent=agent) + after_invocation_event = next(events) + assert isinstance(after_invocation_event, AfterInvocationEvent) + assert after_invocation_event.agent == agent + assert after_invocation_event.result is not None assert len(agent.messages) == 4 @@ -261,7 +264,10 @@ async def test_agent_stream_async_hooks(agent, hook_provider, agent_tool, mock_m ) assert next(events) == MessageAddedEvent(agent=agent, message=agent.messages[3]) - assert next(events) == AfterInvocationEvent(agent=agent) + after_invocation_event = next(events) + assert isinstance(after_invocation_event, AfterInvocationEvent) + assert after_invocation_event.agent == agent + assert after_invocation_event.result is not None assert len(agent.messages) == 4