# Reset sessions
self._agent_thread = self._agent.create_session()

import asyncio
from typing import Any

from agent_framework import Agent, BaseChatClient, ChatResponse, Message
from agent_framework.orchestrations import (
    MagenticBuilder,
    MagenticContext,
    MagenticManagerBase,
    MagenticProgressLedger,
    MagenticProgressLedgerItem,
)


class RecordingClient(BaseChatClient):
    """Offline chat client that records what the agent sends to the model."""

    def __init__(self) -> None:
        super().__init__()
        self.calls: list[list[str]] = []

    def _inner_get_response(self, *, messages: Any, stream: bool, options: Any, **kwargs: Any) -> Any:
        self.calls.append([m.text for m in messages])
        reply = f"reply-{len(self.calls)}"

        async def _get() -> ChatResponse:
            return ChatResponse(messages=Message("assistant", [reply]))

        return _get()


class StallOnceManager(MagenticManagerBase):
    """Asks agentA, stalls once (forcing a reset and replan), asks agentA again, then finishes."""

    def __init__(self) -> None:
        super().__init__(max_round_count=10, max_stall_count=0)
        self.rounds = 0

    async def plan(self, magentic_context: MagenticContext) -> Message:
        return Message("assistant", ["plan"])

    async def replan(self, magentic_context: MagenticContext) -> Message:
        return Message("assistant", ["new plan"])

    async def create_progress_ledger(self, magentic_context: MagenticContext) -> MagenticProgressLedger:
        self.rounds += 1
        progress = self.rounds != 2  # round 2 stalls, which resets and replans
        item = MagenticProgressLedgerItem
        return MagenticProgressLedger(
            is_request_satisfied=item(reason="r", answer=self.rounds >= 4),
            is_in_loop=item(reason="r", answer=not progress),
            is_progress_being_made=item(reason="r", answer=progress),
            next_speaker=item(reason="r", answer="agentA"),
            instruction_or_question=item(reason="r", answer="before reset" if self.rounds == 1 else "after reset"),
        )

    async def prepare_final_answer(self, magentic_context: MagenticContext) -> Message:
        return Message("assistant", ["final"])


async def main() -> None:
    client = RecordingClient()
    agent = Agent(client=client, name="agentA", description="worker")
    workflow = MagenticBuilder(participants=[agent], manager=StallOnceManager()).build()
    await workflow.run("do the task")
    for number, call in enumerate(client.calls, 1):
        print(f"agentA model call {number}: {call}")


asyncio.run(main())