From 8db2dbf3a655df8fb73ea557e2238c7a1eba41e6 Mon Sep 17 00:00:00 2001 From: Gaurav Gupta Date: Fri, 23 Jan 2026 07:53:56 +0000 Subject: [PATCH 1/2] feat: adding terminate attribute for after-model-call-event hooks --- src/strands/event_loop/event_loop.py | 17 ++++ src/strands/hooks/events.py | 6 +- tests/strands/agent/test_agent_hooks.py | 101 ++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 1 deletion(-) diff --git a/src/strands/event_loop/event_loop.py b/src/strands/event_loop/event_loop.py index 9fe645f80..bc50d7a5f 100644 --- a/src/strands/event_loop/event_loop.py +++ b/src/strands/event_loop/event_loop.py @@ -153,6 +153,11 @@ async def event_loop_cycle( agent, cycle_span, cycle_trace, invocation_state, tracer, structured_output_context ) async for model_event in model_events: + if isinstance(model_event, EventLoopStopEvent): + agent.event_loop_metrics.end_cycle(cycle_start_time, cycle_trace) + yield model_event + await model_events.aclose() # clean-up async for-loop to avoid CancelledError + return if not isinstance(model_event, ModelStopReason): yield model_event @@ -360,6 +365,18 @@ async def _handle_model_execution( stop_reason, ) continue # Retry the model call + elif after_model_call_event.terminate: + logger.debug( + "stop_reason=<%s>, termination_requested= | hook requested agent termination", + stop_reason, + ) + invocation_state["request_state"]["stop_event_loop"] = True + yield EventLoopStopEvent( + stop_reason, + message, + agent.event_loop_metrics, + invocation_state["request_state"], + ) if stop_reason == "max_tokens": message = recover_message_on_max_tokens_reached(message) diff --git a/src/strands/hooks/events.py b/src/strands/hooks/events.py index 8d3e5d280..fbb0ab449 100644 --- a/src/strands/hooks/events.py +++ b/src/strands/hooks/events.py @@ -278,9 +278,13 @@ class ModelStopResponse: stop_response: ModelStopResponse | None = None exception: Exception | None = None retry: bool = False + terminate: bool = False def _can_write(self, name: str) -> bool: - return name == "retry" + return name in ( + "retry", + "terminate", + ) @property def should_reverse_callbacks(self) -> bool: diff --git a/tests/strands/agent/test_agent_hooks.py b/tests/strands/agent/test_agent_hooks.py index 4397b9628..00e552afe 100644 --- a/tests/strands/agent/test_agent_hooks.py +++ b/tests/strands/agent/test_agent_hooks.py @@ -694,3 +694,104 @@ async def capture_messages_hook(event: BeforeInvocationEvent): # structured_output_async uses deprecated path that doesn't pass messages assert received_messages is None + + +@pytest.mark.asyncio +async def test_hook_terminate_on_successful_call(): + """Test that hooks can terminate even on successful model calls based on response content.""" + + mock_provider = MockedModelProvider( + [ + { + "role": "assistant", + "content": [{"text": "First conversation successful"}], + }, + { + "role": "assistant", + "content": [{"text": "Unnecessary follow-up conversation"}], + }, + ] + ) + + # Hook that terminate if response is favorable + class SuccessfulTerminateHook: + def __init__(self, end_marker="success"): + self.end_marker = end_marker + self.call_count = 0 + + def register_hooks(self, registry): + registry.add_callback(strands.hooks.AfterModelCallEvent, self.handle_after_model_call) + + async def handle_after_model_call(self, event): + self.call_count += 1 + + # Check successful responses for favorable markers + if event.stop_response: + message = event.stop_response.message + text_content = "".join(block.get("text", "") for block in message.get("content", [])) + + if self.end_marker in text_content: + event.terminate = True + + terminate_hook = SuccessfulTerminateHook(end_marker="success") + agent = Agent(model=mock_provider, hooks=[terminate_hook]) + + result = agent("Generate a response") + + # Verify hook was called only once (For first favorable response) + assert terminate_hook.call_count == 1 + + # Verify final result is the favorable response + assert result.message["content"][0]["text"] == "First conversation successful" + + +@pytest.mark.asyncio +async def test_hook_terminate_gracefully_on_limits(agent_tool, tool_use): + """Test that hooks can terminate agent gracefully after maximum counts reached.""" + + mock_provider = MockedModelProvider( + [ + { + "role": "assistant", + "content": [{"text": "First tool-use"}, {"toolUse": tool_use}], + }, + { + "role": "assistant", + "content": [{"text": "Second tool-use"}, {"toolUse": tool_use}], + }, + { + "role": "assistant", + "content": [{"text": "Third tool-use"}, {"toolUse": tool_use}], + }, + ] + ) + + # Hook that counts number of calls + class GracefulTerminateHook: + def __init__(self, max_counts): + self.max_counts = max_counts + self.call_count = 0 + + def register_hooks(self, registry): + registry.add_callback(strands.hooks.AfterModelCallEvent, self.handle_after_model_call) + + async def handle_after_model_call(self, event): + self.call_count += 1 + + if self.call_count > self.max_counts - 1: + event.terminate = True + + terminate_hook = GracefulTerminateHook(max_counts=2) + agent = Agent( + model=mock_provider, + tools=[agent_tool], + hooks=[terminate_hook], + ) + + result = agent("Generate a response") + + # Verify hook was called two times + assert terminate_hook.call_count == 2 + + # Verify final result is the second tool-use + assert result.message["content"][0]["text"] == "Second tool-use" From ac678a96a634ca390f4eeb96b69af4e7563dca1e Mon Sep 17 00:00:00 2001 From: Gaurav Gupta Date: Wed, 4 Feb 2026 07:47:03 +0000 Subject: [PATCH 2/2] refactor: changing attr name from terminate to stop_loop for consistency --- src/strands/event_loop/event_loop.py | 4 ++-- src/strands/hooks/events.py | 7 +++++-- tests/strands/agent/test_agent_hooks.py | 22 +++++++++++----------- 3 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/strands/event_loop/event_loop.py b/src/strands/event_loop/event_loop.py index bc50d7a5f..990f5fb9f 100644 --- a/src/strands/event_loop/event_loop.py +++ b/src/strands/event_loop/event_loop.py @@ -365,9 +365,9 @@ async def _handle_model_execution( stop_reason, ) continue # Retry the model call - elif after_model_call_event.terminate: + elif after_model_call_event.stop_loop: logger.debug( - "stop_reason=<%s>, termination_requested= | hook requested agent termination", + "stop_reason=<%s>, stop_loop_requested= | hook requested agent stop-loop", stop_reason, ) invocation_state["request_state"]["stop_event_loop"] = True diff --git a/src/strands/hooks/events.py b/src/strands/hooks/events.py index fbb0ab449..ee2d27947 100644 --- a/src/strands/hooks/events.py +++ b/src/strands/hooks/events.py @@ -188,6 +188,9 @@ class AfterToolCallEvent(HookEvent): retry: Whether to retry the tool invocation. Can be set by hook callbacks to trigger a retry. When True, the current result is discarded and the tool is called again. Defaults to False. + stop_loop: Whether to end the event-loop. Hooks can use this flag to terminate + the event-loop immediately. Setting to True would close the event loop and + perform proper closure of async loop. Defaults to False """ selected_tool: AgentTool | None @@ -278,12 +281,12 @@ class ModelStopResponse: stop_response: ModelStopResponse | None = None exception: Exception | None = None retry: bool = False - terminate: bool = False + stop_loop: bool = False def _can_write(self, name: str) -> bool: return name in ( "retry", - "terminate", + "stop_loop", ) @property diff --git a/tests/strands/agent/test_agent_hooks.py b/tests/strands/agent/test_agent_hooks.py index 00e552afe..d2045ca53 100644 --- a/tests/strands/agent/test_agent_hooks.py +++ b/tests/strands/agent/test_agent_hooks.py @@ -697,8 +697,8 @@ async def capture_messages_hook(event: BeforeInvocationEvent): @pytest.mark.asyncio -async def test_hook_terminate_on_successful_call(): - """Test that hooks can terminate even on successful model calls based on response content.""" +async def test_hook_stop_loop_on_successful_call(): + """Test that hooks can stop event loop even on successful model calls based on response content.""" mock_provider = MockedModelProvider( [ @@ -713,8 +713,8 @@ async def test_hook_terminate_on_successful_call(): ] ) - # Hook that terminate if response is favorable - class SuccessfulTerminateHook: + # Hook that stop event loop if response is favorable + class SuccessfulStopLoopHook: def __init__(self, end_marker="success"): self.end_marker = end_marker self.call_count = 0 @@ -731,9 +731,9 @@ async def handle_after_model_call(self, event): text_content = "".join(block.get("text", "") for block in message.get("content", [])) if self.end_marker in text_content: - event.terminate = True + event.stop_loop = True - terminate_hook = SuccessfulTerminateHook(end_marker="success") + terminate_hook = SuccessfulStopLoopHook(end_marker="success") agent = Agent(model=mock_provider, hooks=[terminate_hook]) result = agent("Generate a response") @@ -746,8 +746,8 @@ async def handle_after_model_call(self, event): @pytest.mark.asyncio -async def test_hook_terminate_gracefully_on_limits(agent_tool, tool_use): - """Test that hooks can terminate agent gracefully after maximum counts reached.""" +async def test_hook_stop_loop_gracefully_on_limits(agent_tool, tool_use): + """Test that hooks can stop event-loop of agent gracefully after maximum counts reached.""" mock_provider = MockedModelProvider( [ @@ -767,7 +767,7 @@ async def test_hook_terminate_gracefully_on_limits(agent_tool, tool_use): ) # Hook that counts number of calls - class GracefulTerminateHook: + class GracefulStopLoopHook: def __init__(self, max_counts): self.max_counts = max_counts self.call_count = 0 @@ -779,9 +779,9 @@ async def handle_after_model_call(self, event): self.call_count += 1 if self.call_count > self.max_counts - 1: - event.terminate = True + event.stop_loop = True - terminate_hook = GracefulTerminateHook(max_counts=2) + terminate_hook = GracefulStopLoopHook(max_counts=2) agent = Agent( model=mock_provider, tools=[agent_tool],