Async Errors, Cancellation and Exception Groups

async def generate_and_store(prompt: str, conv_id: str) -> str:
    parts = []
    try:
        async for delta in llm.stream(prompt):
            parts.append(delta)
    except asyncio.CancelledError:
        # user hit "stop": persist the partial answer, then let cancellation continue
        await asyncio.shield(store.save_partial(conv_id, "".join(parts)))
        raise
    text = "".join(parts)
    await store.save(conv_id, text)
    return text

# 3.11+
async def chat_with_deadline(messages: list, seconds: float = 30.0) -> str:
    try:
        async with asyncio.timeout(seconds):
            return await llm.chat(messages)
    except TimeoutError as err:
        raise ProviderTimeoutError(f"no response in {seconds}s", provider="openai", timeout=seconds) from err

async def stream_with_idle_timeout(agen, idle: float = 15.0):
    it = agen.__aiter__()
    while True:
        try:
            chunk = await asyncio.wait_for(it.__anext__(), idle)
        except StopAsyncIteration:
            return
        except asyncio.TimeoutError as err:
            raise StreamInterruptedError(f"no tokens for {idle}s", provider="x") from err
        yield chunk

async def embed_corpus(batches: list) -> list:
    async with asyncio.TaskGroup() as tg:                       # 3.11+
        tasks = [tg.create_task(embedder.embed(b)) for b in batches]
    return [vec for t in tasks for vec in t.result()]

async def embed_one_batch(i: int, batch: list, results: dict, failures: dict) -> None:
    try:
        results[i] = await embedder.embed(batch)
    except ProviderError as err:            # expected per-batch failure: record, do not cancel siblings
        failures[i] = err


async def embed_tolerant(batches: list):
    results, failures = {}, {}
    async with asyncio.TaskGroup() as tg:
        for i, b in enumerate(batches):
            tg.create_task(embed_one_batch(i, b, results, failures))
    return results, failures

try:
    async with asyncio.TaskGroup() as tg:
        tg.create_task(search_vectors(query))
        tg.create_task(search_keywords(query))
        tg.create_task(fetch_user_profile(user_id))
except* VectorStoreError as eg:
    logger.warning("vector search failed: %s", [str(e) for e in eg.exceptions])
    raise RetrievalError("retrieval degraded") from eg
except* (TimeoutError, ProviderTimeoutError) as eg:
    raise ProviderTimeoutError(f"{len(eg.exceptions)} lookups timed out", provider="mixed") from eg

eg = ExceptionGroup("batch failed", [RateLimitError("429", provider="a"), ValueError("bad chunk")])

rate_limited, rest = eg.split(RateLimitError)      # (matching group or None, rest or None)
only_values = eg.subgroup(ValueError)              # matching group or None


def iter_leaves(group):
    """Flatten nested exception groups into leaf exceptions."""
    for exc in group.exceptions:
        if isinstance(exc, BaseExceptionGroup):
            yield from iter_leaves(exc)
        else:
            yield exc


leaves = list(iter_leaves(eg))

def validate_tools(tools: list) -> None:
    errors = []
    for t in tools:
        try:
            validate_tool_schema(t)
        except ValueError as err:
            err.add_note(f"tool={t.get('name')!r}")
            errors.append(err)
    if errors:
        raise ExceptionGroup(f"{len(errors)} invalid tool schemas", errors)

results = await asyncio.gather(*(classify(doc) for doc in docs), return_exceptions=True)
ok = [r for r in results if not isinstance(r, BaseException)]
failed = [(d, r) for d, r in zip(docs, results) if isinstance(r, BaseException)]
for doc, err in failed:
    logger.warning("classify failed for %s", doc.id, exc_info=err)

async def summarise_all(docs: list, concurrency: int = 8) -> dict:
    sem = asyncio.Semaphore(concurrency)
    results: dict = {}
    errors: dict = {}

    async def one(doc) -> None:
        async with sem:
            try:
                results[doc.id] = await with_retry(lambda: llm.summarise(doc.text))
            except (ProviderError, StructuredOutputError) as err:
                errors[doc.id] = err

    async with asyncio.TaskGroup() as tg:
        for d in docs:
            tg.create_task(one(d))

    if errors and not results:
        raise ExceptionGroup("every summary failed", list(errors.values()))
    return {"results": results, "failed": {k: repr(v) for k, v in errors.items()}}

from contextlib import aclosing


async def first_sentence(prompt: str) -> str:
    buf = ""
    async with aclosing(llm.stream(prompt)) as stream:
        async for delta in stream:
            buf += delta
            if "." in buf:
                break                 # stream is closed now, HTTP connection released
    return buf.split(".")[0] + "."

async def sse_events(prompt: str):
    try:
        async for delta in llm.stream(prompt):
            yield f"data: {json.dumps({'delta': delta})}\n\n"
    except ProviderError as err:
        logger.exception("stream failed")
        yield f"event: error\ndata: {json.dumps({'error': err.code, 'retryable': err.retryable})}\n\n"
    yield "data: [DONE]\n\n"

background: set = set()


def spawn(coro, *, name: str) -> asyncio.Task:
    task = asyncio.create_task(coro, name=name)
    background.add(task)                       # keep a strong reference: tasks can be GC'd otherwise
    task.add_done_callback(_on_done)
    return task


def _on_done(task: asyncio.Task) -> None:
    background.discard(task)
    if task.cancelled():
        return
    exc = task.exception()
    if exc is not None:
        logger.error("background task %s failed", task.get_name(), exc_info=exc)

def loop_exception_handler(loop, context):
    exc = context.get("exception")
    logger.error("asyncio: %s", context.get("message"), exc_info=exc)


asyncio.get_running_loop().set_exception_handler(loop_exception_handler)

from concurrent.futures import ThreadPoolExecutor, as_completed

with ThreadPoolExecutor(max_workers=8) as pool:
    futures = {pool.submit(ocr_page, p): p for p in pages}
    for fut in as_completed(futures):
        page = futures[fut]
        try:
            text = fut.result()          # re-raises the worker's exception here
        except OCRError as err:
            logger.warning("page %s failed: %s", page, err)

async def main() -> None:
    server = await start_server()
    try:
        await server.serve_forever()
    finally:
        await vector_store.close()
        await llm_client.aclose()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        sys.exit(130)   # asyncio.run already cancelled main(); finally blocks ran

try:
    from builtins import ExceptionGroup  # 3.11+
except ImportError:  # pragma: no cover
    from exceptiongroup import ExceptionGroup  # type: ignore[no-redef]