Several clients on one session: live views and approvals

import asyncio
from collections.abc import Callable
from dataclasses import dataclass, field

from mistralai import agents

@agents.tool
def publish_report(title: str) -> str:
    """Publish a report to the team's wiki."""
    return f"Published {title!r}."

agent = agents.Agent(
    model="mistral-large-4",
    name="several-clients",
    instructions="You help a team publish reports. Be brief.",
    tools=[publish_report],
    permission_mode="ask",  # publish_report waits for a click
)

@dataclass
class Screen:
    """What one view shows. The app keeps it, so a reload can continue from the cursor."""

    name: str
    cursor: str | None = None
    buttons: dict[str, agents.OpenCallback] = field(default_factory=dict)
    lines: list[str] = field(default_factory=list)
    shown: str = ""

    def update(self) -> None:
        """Redraw, here as one printed line when something changed."""
        now = f"buttons: {[b.title for b in self.buttons.values()]}, answers: {len(self.lines)}"
        if now != self.shown:
            print(f"  [{self.name}] {now}")
            self.shown = now

async def view(session: agents.Session, screen: Screen) -> None:
    """One view (desktop or web). A reload continues from screen.cursor."""
    async for event in session.iter_events(after=screen.cursor):
        match event:
            case agents.CallbackOpened(callback=callback):
                screen.buttons[callback.id] = callback
            case agents.CallbackResolved(callback_id=callback_id):
                screen.buttons.pop(callback_id, None)
            case agents.TurnEnded(state=state):
                screen.lines.append(state.text)
        screen.cursor = event.id
        screen.update()

async def click(
    session: agents.Session, callback: agents.OpenCallback, mine: agents.AllowApproval | agents.DenyApproval
) -> str:
    match await session.turn(callback.turn_id).answer(callback.id, mine):
        case agents.Answered(answer=applied) if applied != mine:
            return "the other view answered first"
        case agents.Answered():
            return "done"
        case agents.AlreadySettled():
            return "someone else answered"
        case agents.NotFound():
            return "no such approval"

async def on_typed(session: agents.Session, text: str) -> None:
    await session.send_message(text)  # starts a turn, or steers the running one

async def until(check: Callable[[], bool], seconds: float = 120) -> None:
    async with asyncio.timeout(seconds):  # a wait that never ends is a bug: fail instead
        while not check():
            await asyncio.sleep(0.05)

async def main() -> tuple[Screen, Screen]:
    session = await agent.start_session()
    desktop, web = Screen("desktop"), Screen("web")
    views = {screen.name: asyncio.create_task(view(session, screen)) for screen in (desktop, web)}
    try:
        print("> Publish the weekly report, titled 'Week 44'.")
        await on_typed(session, "Publish the weekly report, titled 'Week 44'.")
        await until(lambda: bool(desktop.buttons) and bool(web.buttons))

        # Both views click at once: the desktop allows, the web view denies.
        callback = next(iter(desktop.buttons.values()))
        allow, deny = await asyncio.gather(
            click(session, callback, agents.AllowApproval()),
            click(session, callback, agents.DenyApproval()),
        )
        print(f"desktop clicked Allow: {allow}; web clicked Deny: {deny}")

        # When the Deny counts, the model may ask again: the desktop allows each new approval.
        answered = {callback.id}

        async def allow_new_approvals() -> None:
            for again in [b for b in desktop.buttons.values() if b.id not in answered]:
                answered.add(again.id)
                print(f"desktop clicked Allow again: {await click(session, again, agents.AllowApproval())}")

        async with asyncio.timeout(120):
            while not (len(desktop.lines) == 1 and len(web.lines) == 1):
                await allow_new_approvals()
                await asyncio.sleep(0.05)

        # The web view reloads: a new stream from its cursor, so no answer shows twice.
        views["web"].cancel()
        views["web"] = asyncio.create_task(view(session, web))
        print("> Thanks!")
        await on_typed(session, "Thanks!")
        await until(lambda: len(desktop.lines) == 2 and len(web.lines) == 2)
        for line in web.lines:
            print(f"web: {line}")
        return desktop, web
    finally:
        for task in views.values():
            task.cancel()
        await session.cancel()

if __name__ == "__main__":
    asyncio.run(main())

"""The test of the several clients cookbook: a scripted model, no API key.

Run it beside several_clients.py: pytest -o asyncio_mode=auto test_several_clients.py
"""

import asyncio

import pytest
import several_clients

from mistralai import agents

async def test_one_click_counts_and_a_reload_shows_no_repeats(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    model = agents.testing.ScriptedModel(
        [
            # The model calls the tool from code, as a real model does.
            agents.testing.run_typescript('return await tools.client.publish_report({"title": "Week 44"});'),
            agents.testing.reply("Published 'Week 44'."),
            agents.testing.reply("You're welcome."),
        ]
    )
    monkeypatch.setattr(several_clients, "agent", several_clients.agent.override(model=model))

    desktop, web = await several_clients.main()

    model.assert_consumed()
    assert desktop.lines == web.lines == ["Published 'Week 44'.", "You're welcome."]
    assert "desktop clicked Allow: done; web clicked Deny: the other view answered first" in capsys.readouterr().out

async def test_when_the_deny_counts_the_desktop_allows_the_next_approval(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    model = agents.testing.ScriptedModel(
        [
            agents.testing.run_typescript('return await tools.client.publish_report({"title": "Week 44"});'),
            # The web view's Deny counted, so the model asks again.
            agents.testing.run_typescript('return await tools.client.publish_report({"title": "Week 44"});'),
            agents.testing.reply("Published 'Week 44'."),
            agents.testing.reply("You're welcome."),
        ]
    )
    monkeypatch.setattr(several_clients, "agent", several_clients.agent.override(model=model))
    click, denied = several_clients.click, asyncio.Event()

    async def deny_first(session, callback, mine):  # the Allow waits until the Deny has counted
        if isinstance(mine, agents.AllowApproval) and not denied.is_set():
            await denied.wait()
        result = await click(session, callback, mine)
        if isinstance(mine, agents.DenyApproval):
            denied.set()
        return result

    monkeypatch.setattr(several_clients, "click", deny_first)

    desktop, web = await several_clients.main()

    model.assert_consumed()
    assert desktop.lines == web.lines == ["Published 'Week 44'.", "You're welcome."]
    out = capsys.readouterr().out
    assert "desktop clicked Allow: someone else answered; web clicked Deny: done" in out
    assert "desktop clicked Allow again: done" in out