ai-python-custom-loop

Keep the default shape unless you must change control flow:

class MyAgent(ai.Agent):
    async def loop(self, context: ai.Context):
        while context.keep_running():
            async with (
                ai.stream(context=context) as stream,
                ai.ToolRunner() as runner,
            ):
                async for event in ai.util.merge(stream, runner.events()):
                    yield event

                    if isinstance(event, ai.events.ToolEnd):
                        runner.schedule(context.resolve(event.tool_call))