Stop, re-run and answer a turn by id from another kernel

import asyncio
import os
import re
from pathlib import Path

import httpx

from mistralai import agents

REQUIREMENTS = Path("requirements.txt")
PACKAGE = re.compile(r"[A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?")  # a package name, as PEP 508 spells it
VERSION = re.compile(r"[0-9]+(\.[0-9]+)*")

@agents.tool
async def release_notes(repository: str) -> str:
    """The notes of a GitHub repository's last 6 releases, such as encode/httpx."""
    # GitHub rate-limits calls without a token: set GITHUB_TOKEN to raise the limit.
    token = os.environ.get("GITHUB_TOKEN")
    async with httpx.AsyncClient(timeout=10, headers={"Authorization": f"Bearer {token}"} if token else {}) as client:
        response = await client.get(f"https://api.github.com/repos/{repository}/releases", params={"per_page": 6})
        response.raise_for_status()
    return "\n\n".join(f"## {release['tag_name']}\n{release['body']}" for release in response.json())

@agents.tool
async def latest_release(package: str) -> str:
    """A package's latest version on PyPI."""
    async with httpx.AsyncClient(timeout=10) as client:
        response = await client.get(f"https://pypi.org/pypi/{package}/json")
        response.raise_for_status()
    return f"{package} {response.json()['info']['version']}"

@agents.tool
async def versions(package: str) -> str:
    """A package's last 10 final releases on PyPI (such as 1.2.3), oldest first, without yanked ones."""
    async with httpx.AsyncClient(timeout=10) as client:
        response = await client.get(f"https://pypi.org/pypi/{package}/json")
        response.raise_for_status()
    final = [
        version
        for version, files in response.json()["releases"].items()
        if version.replace(".", "").isdigit() and files and not all(file["yanked"] for file in files)
    ]
    return ", ".join(sorted(final, key=lambda version: tuple(map(int, version.split("."))))[-10:])

def name(line: str) -> str:
    """The package a requirements.txt line names, as pip compares names: httpx for HTTPX[http2]>=0.27."""
    return re.sub(r"[-_.]+", "-", re.split(r"[^A-Za-z0-9._-]", line.strip(), maxsplit=1)[0]).lower()

@agents.tool
def pin(package: str, version: str) -> str:
    """Pin a package to a version in requirements.txt."""
    # The model's text goes into a file that pip reads: a package name and a version, nothing else.
    if not PACKAGE.fullmatch(package) or not VERSION.fullmatch(version):
        raise agents.ToolError(f"Not a package name and a version: {package!r} {version!r}.")
    lines = REQUIREMENTS.read_text().splitlines()
    mine = [line for line in lines if name(line) == name(package)]
    # Extras, markers, ranges and comments on the package's line are a person's to rewrite.
    if any(not re.fullmatch(r"[A-Za-z0-9._-]+==[0-9.]+", line) for line in mine):
        raise agents.ToolError(f"requirements.txt has more than a version for {package}: edit its line by hand.")
    new = f"{package}=={version}"
    lines = [new if line in mine else line for line in lines] if mine else [*lines, new]
    REQUIREMENTS.write_text("".join(f"{line}\n" for line in lines))
    return f"requirements.txt pins {new}"

reader = agents.Agent(
    model="mistral-large-4",
    name="release-notes-reader",
    instructions=(
        "You review dependency upgrades for a Python service. Read the package's release notes with "
        "release_notes, and list each change that can break code that calls the package, one line "
        "each, with the version that made it and what to use instead. Say only what the notes say. At "
        "most 6 lines, no preamble."
    ),
    tools=[release_notes],
)
upgrader = agents.Agent(
    model="mistral-large-4",
    name="upgrader",
    instructions=(
        "You pin the dependencies of a Python service. Find the versions with versions, pin the one "
        "the task asks for with pin. Then answer in one line, without markdown: the new pin, the "
        "one change that makes it the newest one that works, and what the code should move to."
    ),
    tools=[versions, pin],
    # Looking a version up runs; editing requirements.txt waits for an approval.
    permission_policy=agents.PermissionPolicy(
        default="allow",
        rules=(agents.PermissionRule(tools=("client.pin",), decision="ask"),),
    ),
)
reporter = agents.Agent(
    model="mistral-large-4",
    name="dependency-report",
    instructions=(
        "You write the dependency report of a Python service. Look up each pinned package with "
        "latest_release, then write one line per package: the pin, the latest version, and whether "
        "the pin is behind. No preamble."
    ),
    tools=[latest_release],
)

async def kernel_1() -> tuple[str, str]:
    question = "What in httpx's last releases can break our client code? The repository is encode/httpx."
    call = asyncio.create_task(reader.run(question, idempotency_key="httpx-notes-1"))
    await asyncio.sleep(1)
    call.cancel()  # the stop button, after 1 s
    try:
        await call
    except asyncio.CancelledError as error:
        print(f"CancelledError, note: {error.__notes__[0]}")  # the ids, and how to stop the turn

    # The sessions running in this process: the stopped call's session is still one of them.
    print([info.id async for info in agents.list_sessions(host="local", status="running")])
    # The cell again, with the same key: the same turn comes back, not a second one, and finishes.
    state = await reader.run(question, idempotency_key="httpx-notes-1")
    print(f"same key: {type(state).__name__}, turn {state.turn_id}:\n{state.text}")
    await reader.session(state.session_id).cancel()

    task = (
        "Our client code calls httpx.Client(proxies=...). Pin httpx to the newest release that still "
        "has that argument, until we move to its replacement. The breaking changes in httpx's last "
        f"releases:\n{state.text}"
    )
    state = await upgrader.run(task, idempotency_key="pin-httpx-1")
    print(f"pin: {type(state).__name__}")  # Blocked: no on_approval
    return state.session_id, state.turn_id

async def report_cell(session: agents.Session) -> agents.Turn:
    """A notebook cell: start the dependency report, or steer the report that still runs."""
    match await session.start_turn(f"Write the dependency report of these pins:\n{REQUIREMENTS.read_text()}"):
        case agents.Busy(turn=running):  # a session runs one turn at a time
            print("again: Busy")
            await running.steer("End with the number of pins that are behind.")
            return running
        case turn:
            return turn

async def kernel_2(session_id: str, turn_id: str) -> None:
    """Only the pasted ids carry over."""

    def review(ctx: agents.ApprovalCallbackContext, request: agents.ApprovalRequest) -> agents.AllowApproval:
        print(f"approve: {request.tool.id} {request.input}")
        return agents.AllowApproval()

    session = upgrader.session(session_id, on_approval=review)
    state = await session.turn(turn_id).result()
    print(f"pin: {type(state).__name__}: {state.text}")  # Completed: the edit ran once
    await session.cancel()

    session = await reporter.start_session()
    first = await report_cell(session)
    await report_cell(session)  # the cell ran twice
    state = await first.result()
    print(f"report: {type(state).__name__}:\n{state.text}")
    await session.cancel()

async def main() -> None:
    REQUIREMENTS.write_text("httpx==0.27.0\npydantic==2.9.2\nrich==13.9.4\n")  # the service's pins
    session_id, turn_id = await kernel_1()
    print(f"ids to paste: {session_id} {turn_id}")
    await kernel_2(session_id, turn_id)
    print(f"requirements.txt:\n{REQUIREMENTS.read_text()}", end="")

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