From e8df3d8515a372ffdd77b7f8cb9138f76c644822 Mon Sep 17 00:00:00 2001 From: harshdrkm Date: Mon, 2 Nov 2026 12:46:17 -0800 Subject: [PATCH 1/2] feat: add an awaited after-step agent hook Signed-off-by: harshdrkm --- .../content/docs/reference/ai/agent.mdx | 18 ++ src/ai/agents/agent.py | 15 ++ tests/agents/test_after_step.py | 197 ++++++++++++++++++ 3 files changed, 230 insertions(+) create mode 100644 tests/agents/test_after_step.py diff --git a/docs/ai-python/content/docs/reference/ai/agent.mdx b/docs/ai-python/content/docs/reference/ai/agent.mdx index 02692e6e..476f6bd5 100644 --- a/docs/ai-python/content/docs/reference/ai/agent.mdx +++ b/docs/ai-python/content/docs/reference/ai/agent.mdx @@ -107,6 +107,24 @@ stream.output text. When `output_type` is set, the text is validated as JSON against that Pydantic model and the parsed instance is returned. +## after_step + +Override `async def after_step(self, context: ai.Context) -> None` to inspect +or modify the context between steps without replacing the default loop. +The default implementation does nothing. + +The default loop calls it once after each completed model response and all +of its tool results have been appended to history, after the step's stream, +tool runner, and event merge have closed. This includes a final response +without tool calls. Changes to the context affect the next `keep_running()` +check; for example, appending a pending user message continues the run. +An exception or cancellation propagates and prevents the next model request. +A failed step does not call `after_step`. + +The hook does not wait for the consumer to drain buffered events. Set +`LOOP_BUFFER = 0` if it must run after consumer side effects for the step's +events. A custom `loop` must call `after_step` itself if needed. + ## loop Override `Agent.loop(context)` to customize control flow. The default loop uses diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index ec53583c..fe77aeb2 100644 --- a/src/ai/agents/agent.py +++ b/src/ai/agents/agent.py @@ -1460,6 +1460,19 @@ def resolve( """Resolve ToolCallPart(s) into callable ToolCall object(s).""" return self._tool_resolver.resolve(tool_part) + async def after_step(self, context: Context) -> None: + """Run after a completed step, before deciding whether to continue. + + The assistant message and all tool results are in ``context``, and + the step's stream and tool runner are closed. This includes final + assistant-only steps. Override to inspect or modify the context; + the default does nothing. Exceptions and cancellation propagate. + + This does not wait for the consumer to drain buffered events. Use + ``LOOP_BUFFER = 0`` when sequencing against consumer side effects. + Custom ``loop`` implementations must call this method themselves. + """ + async def loop( self, context: Context ) -> AsyncGenerator[events_.AgentEvent]: @@ -1490,6 +1503,8 @@ async def loop( # the loop. context.add(tr.get_tool_message()) + await self.after_step(context) + @overload def run( self, diff --git a/tests/agents/test_after_step.py b/tests/agents/test_after_step.py new file mode 100644 index 00000000..858fbbb1 --- /dev/null +++ b/tests/agents/test_after_step.py @@ -0,0 +1,197 @@ +"""Default-loop step boundaries through the public Agent.run API.""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator +from typing import Any + +import pytest + +import ai +from ai.types import events, messages + +from ..conftest import ( + MOCK_MODEL, + MOCK_PROVIDER, + Recorder, + collect_messages, + mock_llm, + text_msg, + tool_call_msg, +) + + +@pytest.mark.parametrize("buffer", [None, 0]) +async def test_after_step_can_continue_final_response( + buffer: int | None, +) -> None: + snapshots: list[list[str]] = [] + consumed_responses: list[str] = [] + + class PendingInputAgent(ai.Agent): + LOOP_BUFFER = buffer + + async def after_step(self, context: ai.Context) -> None: + snapshots.append([message.text for message in context.messages]) + if buffer == 0: + assert len(consumed_responses) == len(snapshots) + assert consumed_responses[-1] == context.messages[-1].text + if len(snapshots) == 1: + context.add(ai.user_message("One more question")) + + llm = mock_llm([[text_msg("First answer")], [text_msg("Second answer")]]) + async with PendingInputAgent().run( + MOCK_MODEL, [ai.user_message("Question")] + ) as stream: + async for event in stream: + if isinstance(event, events.StreamEnd): + consumed_responses.append(event.message.text) + assert stream.output == "Second answer" + + assert llm.call_count == 2 + assert snapshots == [ + ["Question", "First answer"], + ["Question", "First answer", "One more question", "Second answer"], + ] + + +@pytest.mark.parametrize("first", [1, 2]) +@pytest.mark.parametrize("buffer", [None, 0]) +async def test_after_step_sees_complete_concurrent_batch( + first: int, buffer: int | None, recorder: Recorder +) -> None: + started: list[int] = [] + finished: list[int] = [] + both_started = asyncio.Event() + first_finished = asyncio.Event() + snapshots: list[list[messages.Message]] = [] + consumed_results: set[str] = set() + + @ai.tool + async def complete(number: int) -> int: + """Complete in the specified order after both tools have started.""" + started.append(number) + if len(started) == 2: + both_started.set() + await both_started.wait() + if number != first: + await first_finished.wait() + finished.append(number) + if number == first: + first_finished.set() + return number * 10 + + class BatchAgent(ai.Agent): + LOOP_BUFFER = buffer + + async def after_step(self, context: ai.Context) -> None: + assert finished == [first, 3 - first] + # The outermost step context exits after stream/runner/merge. + assert recorder.ended[-1].name == "loop_turn" + if buffer == 0: + assert consumed_results == {"tc-1", "tc-2"} + snapshots.append(list(context.messages)) + + calls = messages.Message( + role="assistant", + parts=[ + messages.ToolCallPart( + tool_call_id=f"tc-{number}", + tool_name="complete", + tool_args=f'{{"number": {number}}}', + ) + for number in (1, 2) + ], + ) + llm = mock_llm([[calls], [text_msg("Done")]]) + async with asyncio.timeout(5): + async with BatchAgent(tools=[complete]).run( + MOCK_MODEL, [ai.user_message("Run both")] + ) as stream: + async for event in stream: + if isinstance(event, events.ToolCallResult): + consumed_results.update( + result.tool_call_id for result in event.results + ) + + assert llm.call_count == 2 + assert len(snapshots) == 2 + assert [message.role for message in snapshots[0]] == [ + "user", + "assistant", + "tool", + ] + assert len(snapshots[0][1].tool_calls) == 2 + assert { + result.tool_call_id: result.result + for result in snapshots[0][-1].tool_results + } == {"tc-1": 10, "tc-2": 20} + assert snapshots[1][-1].text == "Done" + + +@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError]) +async def test_after_step_failure_stops_before_next_model( + failure: type[RuntimeError] | type[asyncio.CancelledError], +) -> None: + calls = 0 + + @ai.tool + async def complete() -> str: + """Return a result that would ordinarily request another model step.""" + return "result" + + class FailingAgent(ai.Agent): + async def after_step(self, context: ai.Context) -> None: + nonlocal calls + calls += 1 + assert context.messages[-1].tool_results[0].result == "result" + raise failure("after-step failure") + + llm = mock_llm( + [ + [tool_call_msg(name="complete")], + [text_msg("Should not be requested")], + ] + ) + expected = ExceptionGroup if failure is RuntimeError else failure + with pytest.raises(expected) as caught: + async with FailingAgent(tools=[complete]).run( + MOCK_MODEL, [ai.user_message("Run tool")] + ) as stream: + await collect_messages(stream) + + error = caught.value + if isinstance(error, ExceptionGroup): + (error,) = error.exceptions + assert isinstance(error, failure) + assert str(error) == "after-step failure" + assert calls == 1 + assert llm.call_count == 1 + + +async def test_after_step_skips_failed_model_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = 0 + + class ObservingAgent(ai.Agent): + async def after_step(self, context: ai.Context) -> None: + nonlocal calls + calls += 1 + + async def failing_stream( + *args: Any, **kwargs: Any + ) -> AsyncGenerator[events.Event]: + yield events.StreamStart() + raise RuntimeError("provider disconnected") + + monkeypatch.setattr(MOCK_PROVIDER, "_stream_impl", failing_stream) + with pytest.raises(ExceptionGroup, match="unhandled errors") as caught: + async with ObservingAgent().run( + MOCK_MODEL, [ai.user_message("Question")] + ) as stream: + await collect_messages(stream) + + assert "provider disconnected" in str(caught.value.exceptions) + assert calls == 0 From 8c5a9d969f90be5f01779e0321efa73e6b0976f2 Mon Sep 17 00:00:00 2001 From: harshdrkm Date: Mon, 2 Nov 2026 13:15:28 -0800 Subject: [PATCH 2/2] test: focus after-step coverage on continuation and completed batches Signed-off-by: harshdrkm --- tests/agents/test_after_step.py | 143 +++++--------------------------- 1 file changed, 20 insertions(+), 123 deletions(-) diff --git a/tests/agents/test_after_step.py b/tests/agents/test_after_step.py index 858fbbb1..b044ec28 100644 --- a/tests/agents/test_after_step.py +++ b/tests/agents/test_after_step.py @@ -3,18 +3,12 @@ from __future__ import annotations import asyncio -from collections.abc import AsyncGenerator -from typing import Any - -import pytest import ai from ai.types import events, messages from ..conftest import ( MOCK_MODEL, - MOCK_PROVIDER, - Recorder, collect_messages, mock_llm, text_msg, @@ -22,21 +16,12 @@ ) -@pytest.mark.parametrize("buffer", [None, 0]) -async def test_after_step_can_continue_final_response( - buffer: int | None, -) -> None: +async def test_after_step_can_continue_final_response() -> None: snapshots: list[list[str]] = [] - consumed_responses: list[str] = [] class PendingInputAgent(ai.Agent): - LOOP_BUFFER = buffer - async def after_step(self, context: ai.Context) -> None: snapshots.append([message.text for message in context.messages]) - if buffer == 0: - assert len(consumed_responses) == len(snapshots) - assert consumed_responses[-1] == context.messages[-1].text if len(snapshots) == 1: context.add(ai.user_message("One more question")) @@ -44,9 +29,7 @@ async def after_step(self, context: ai.Context) -> None: async with PendingInputAgent().run( MOCK_MODEL, [ai.user_message("Question")] ) as stream: - async for event in stream: - if isinstance(event, events.StreamEnd): - consumed_responses.append(event.message.text) + await collect_messages(stream) assert stream.output == "Second answer" assert llm.call_count == 2 @@ -56,55 +39,42 @@ async def after_step(self, context: ai.Context) -> None: ] -@pytest.mark.parametrize("first", [1, 2]) -@pytest.mark.parametrize("buffer", [None, 0]) -async def test_after_step_sees_complete_concurrent_batch( - first: int, buffer: int | None, recorder: Recorder -) -> None: +async def test_after_step_sees_complete_concurrent_batch() -> None: started: list[int] = [] - finished: list[int] = [] both_started = asyncio.Event() - first_finished = asyncio.Event() + second_finished = asyncio.Event() snapshots: list[list[messages.Message]] = [] consumed_results: set[str] = set() @ai.tool async def complete(number: int) -> int: - """Complete in the specified order after both tools have started.""" + """Start both tools, then complete the second one first.""" started.append(number) if len(started) == 2: both_started.set() await both_started.wait() - if number != first: - await first_finished.wait() - finished.append(number) - if number == first: - first_finished.set() + if number == 1: + await second_finished.wait() + else: + second_finished.set() return number * 10 class BatchAgent(ai.Agent): - LOOP_BUFFER = buffer + LOOP_BUFFER = 0 async def after_step(self, context: ai.Context) -> None: - assert finished == [first, 3 - first] - # The outermost step context exits after stream/runner/merge. - assert recorder.ended[-1].name == "loop_turn" - if buffer == 0: - assert consumed_results == {"tc-1", "tc-2"} + assert consumed_results == {"tc-1", "tc-2"} snapshots.append(list(context.messages)) - calls = messages.Message( - role="assistant", - parts=[ - messages.ToolCallPart( - tool_call_id=f"tc-{number}", - tool_name="complete", - tool_args=f'{{"number": {number}}}', - ) - for number in (1, 2) - ], - ) - llm = mock_llm([[calls], [text_msg("Done")]]) + calls = [ + tool_call_msg( + tc_id=f"tc-{number}", + name="complete", + args=f'{{"number": {number}}}', + ) + for number in (1, 2) + ] + llm = mock_llm([calls, [text_msg("Done")]]) async with asyncio.timeout(5): async with BatchAgent(tools=[complete]).run( MOCK_MODEL, [ai.user_message("Run both")] @@ -117,81 +87,8 @@ async def after_step(self, context: ai.Context) -> None: assert llm.call_count == 2 assert len(snapshots) == 2 - assert [message.role for message in snapshots[0]] == [ - "user", - "assistant", - "tool", - ] - assert len(snapshots[0][1].tool_calls) == 2 assert { result.tool_call_id: result.result for result in snapshots[0][-1].tool_results } == {"tc-1": 10, "tc-2": 20} assert snapshots[1][-1].text == "Done" - - -@pytest.mark.parametrize("failure", [RuntimeError, asyncio.CancelledError]) -async def test_after_step_failure_stops_before_next_model( - failure: type[RuntimeError] | type[asyncio.CancelledError], -) -> None: - calls = 0 - - @ai.tool - async def complete() -> str: - """Return a result that would ordinarily request another model step.""" - return "result" - - class FailingAgent(ai.Agent): - async def after_step(self, context: ai.Context) -> None: - nonlocal calls - calls += 1 - assert context.messages[-1].tool_results[0].result == "result" - raise failure("after-step failure") - - llm = mock_llm( - [ - [tool_call_msg(name="complete")], - [text_msg("Should not be requested")], - ] - ) - expected = ExceptionGroup if failure is RuntimeError else failure - with pytest.raises(expected) as caught: - async with FailingAgent(tools=[complete]).run( - MOCK_MODEL, [ai.user_message("Run tool")] - ) as stream: - await collect_messages(stream) - - error = caught.value - if isinstance(error, ExceptionGroup): - (error,) = error.exceptions - assert isinstance(error, failure) - assert str(error) == "after-step failure" - assert calls == 1 - assert llm.call_count == 1 - - -async def test_after_step_skips_failed_model_stream( - monkeypatch: pytest.MonkeyPatch, -) -> None: - calls = 0 - - class ObservingAgent(ai.Agent): - async def after_step(self, context: ai.Context) -> None: - nonlocal calls - calls += 1 - - async def failing_stream( - *args: Any, **kwargs: Any - ) -> AsyncGenerator[events.Event]: - yield events.StreamStart() - raise RuntimeError("provider disconnected") - - monkeypatch.setattr(MOCK_PROVIDER, "_stream_impl", failing_stream) - with pytest.raises(ExceptionGroup, match="unhandled errors") as caught: - async with ObservingAgent().run( - MOCK_MODEL, [ai.user_message("Question")] - ) as stream: - await collect_messages(stream) - - assert "provider disconnected" in str(caught.value.exceptions) - assert calls == 0