Implementing abort functionality in LangGraph (Python)

[david1542]

Hey everyone, I'm trying to implement an "abort" functionality in my langgraph agent.

More specifically, a user might click "Abort" in my UI. In this case, I'm changing some state in the DB and I'm checking it periodically in parallel to the LangGraph astream loop.
I've tried to create a race condition between these 2 things, but seems like whenever my graph is stuck on some async call (e.g an LLM call), the cancellation doesn't work.

This is my simplified code:

import asyncio
from typing import Callable


async def check_abort(abort_fn: Callable[[], bool]):
    while True:
        await asyncio.sleep(1)  # Check every second
        is_aborted = await abort_fn()
        if is_aborted:
            print("Abort signal detected")
            return


async def execute_graph():
    async for event, state in graph.astream(
        input=input, config=config, stream_mode="values", subgraphs=True
    ):
        print(event, state)


async def is_aborted():
  # Some logic to check against the DB if the run is aborted

async def main():
    try:
        # Create both tasks
        abort_checker = asyncio.create_task(check_abort(is_aborted))
        graph_executor = asyncio.create_task(execute_graph())

        # Race between the entire stream processing and abort checker
        done, pending = await asyncio.wait(
            [abort_checker, graph_executor], return_when=asyncio.FIRST_COMPLETED
        )

        # Cancel whichever task didn't complete
        for task in pending:
            task.cancel()

        if abort_checker in done:
            print("Aborting investigation...")
            print("Cleaning up things...")
            await do_some_cleanup()

        # Wait for cancelled tasks to finish cleanup
        await asyncio.gather(*pending, return_exceptions=True)

    finally:
        # Ensure both tasks are cleaned up
        for task in [abort_checker, graph_executor]:
            if not task.done():
                task.cancel()


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

As I said, when the graph is stuck on an LLM call, the race condition doesn't work and the execution proceeds for some reason. I'm not sure why :face_with_monocle:

Does any of you implemented an abort functionality maybe? I'd love to get some advice.

[hinthornw]

What does the node with the LLM call look like?