Add an awaited after_step extension to the default agent loop

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

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,

tests/agents/test_after_step.py

@@ -0,0 +1,94 @@
+"""Default-loop step boundaries through the public Agent.run API."""
+
+from __future__ import annotations
+
+import asyncio
+
+import ai
+from ai.types import events, messages
+
+from ..conftest import (
+    MOCK_MODEL,
+    collect_messages,
+    mock_llm,
+    text_msg,
+    tool_call_msg,
+)
+
+
+async def test_after_step_can_continue_final_response() -> None:
+    snapshots: list[list[str]] = []
+
+    class PendingInputAgent(ai.Agent):
+        async def after_step(self, context: ai.Context) -> None:
+            snapshots.append([message.text for message in context.messages])
+            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:
+        await collect_messages(stream)
+        assert stream.output == "Second answer"
+
+    assert llm.call_count == 2
+    assert snapshots == [
+        ["Question", "First answer"],
+        ["Question", "First answer", "One more question", "Second answer"],
+    ]
+
+
+async def test_after_step_sees_complete_concurrent_batch() -> None:
+    started: list[int] = []
+    both_started = asyncio.Event()
+    second_finished = asyncio.Event()
+    snapshots: list[list[messages.Message]] = []
+    consumed_results: set[str] = set()
+
+    @ai.tool
+    async def complete(number: int) -> int:
+        """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 == 1:
+            await second_finished.wait()
+        else:
+            second_finished.set()
+        return number * 10
+
+    class BatchAgent(ai.Agent):
+        LOOP_BUFFER = 0
+
+        async def after_step(self, context: ai.Context) -> None:
+            assert consumed_results == {"tc-1", "tc-2"}
+            snapshots.append(list(context.messages))
+
+    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")]
+        ) 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 {
+        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"