Try two coding fixes and keep the passing branch

import asyncio
from contree_sdk import Contree

TEST_SOURCE = """
from calc import add

assert add(2, 3) == 5
assert add(-1, 1) == 0
print("PASS")
"""
BROKEN_SOURCE = """
def add(a, b):
    return a - b
"""
MULTIPLY_SOURCE = """
def add(a, b):
    return a * b
"""
ADDITION_SOURCE = """
def add(a, b):
    return a + b
"""
FOLLOW_UP_SOURCE = """
from calc import add

assert add(10, 7) == 17
print("FOLLOW-UP PASS")
"""

async def main() -> None:
    client = Contree()
    base = await client.images.oci("docker.io/library/python:3.12-alpine")
    prepared = await base.run(
        shell="true",
        files={
            "/repo/calc.py": BROKEN_SOURCE.encode(),
            "/repo/test_calc.py": TEST_SOURCE.encode(),
        },
        disposable=False,
        timeout=60,
    )
    if prepared.exit_code != 0 or prepared.uuid is None:
        raise RuntimeError('Preparation failed or returned no saved image')
    print(f"Prepared test image: {prepared.uuid}")
    baseline = await client.images.use(prepared.uuid, strict=True)
    baseline_test = await baseline.run(
        shell="python test_calc.py",
        cwd="/repo",
        disposable=True,
        timeout=30,
    )
    if baseline_test.exit_code == 0:
        raise RuntimeError('The baseline should fail before either patch is applied')

    print("Baseline fails as expected; trying both fixes")
    multiply_result, addition_result = await asyncio.gather(
        baseline.run(
            shell="python test_calc.py",
            files={"/repo/calc.py": MULTIPLY_SOURCE.encode()},
            cwd="/repo",
            disposable=False,
            timeout=30,
        ),
        baseline.run(
            shell="python test_calc.py",
            files={"/repo/calc.py": ADDITION_SOURCE.encode()},
            cwd="/repo",
            disposable=False,
            timeout=30,
        ),
    )
    outcomes = {"multiply": multiply_result.exit_code, "add": addition_result.exit_code}
    if outcomes != {"multiply": 1, "add": 0}:
        raise RuntimeError('Expected the multiplication patch to fail and the addition patch to pass')
    if addition_result.stdout != "PASS\n" or addition_result.uuid is None:
        raise RuntimeError('Passing patch did not return PASS output and a saved image')
    selected_source = await addition_result.read("/repo/calc.py")
    if selected_source != ADDITION_SOURCE.encode():
        raise RuntimeError('Saved source does not contain the selected addition fix')

    print(f"Patch results: {outcomes}; continuing from {addition_result.uuid}")
    selected_image = await client.images.use(addition_result.uuid, strict=True)
    continued = await selected_image.run(
        shell="python test_calc.py && python follow_up.py",
        files={"/repo/follow_up.py": FOLLOW_UP_SOURCE.encode()},
        cwd="/repo",
        disposable=True,
        timeout=30,
    )
    if continued.exit_code != 0:
        raise RuntimeError('The selected image failed the follow-up checks')
    print(continued.stdout, end="")
    print(f"baseline={prepared.uuid} chosen={addition_result.uuid} outcomes={outcomes}")


asyncio.run(main())