From feb4e908ae8d2a4d27bde4e29dbfaf0ed71d9593 Mon Sep 17 00:00:00 2001
From: peterychang <49209570+peterychang@users.noreply.github.com>
Date: Wed, 20 Aug 2025 15:56:26 -0400
Subject: [PATCH] Python: Request cancellation sample (#459)

* request cancellation via tasks

* fix missing kwargs
---
 .../agent_framework/_cancellation_token.py    | 50 -------------------
 .../packages/main/agent_framework/_clients.py |  2 +-
 .../chat_client/chat_response_cancellation.py | 19 +++++++
 3 files changed, 20 insertions(+), 51 deletions(-)
 delete mode 100644 python/packages/main/agent_framework/_cancellation_token.py
 create mode 100644 python/samples/getting_started/chat_client/chat_response_cancellation.py

diff --git a/python/packages/main/agent_framework/_cancellation_token.py b/python/packages/main/agent_framework/_cancellation_token.py
deleted file mode 100644
index 770200fbbd..0000000000
--- a/python/packages/main/agent_framework/_cancellation_token.py
+++ /dev/null
@@ -1,50 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-import threading
-from asyncio import Future
-from collections.abc import Callable
-from typing import Any
-
-
-# from https://github.com/microsoft/autogen/blob/main/python/packages/autogen-core/src/autogen_core/_cancellation_token.py
-class CancellationToken:
-    """A token used to cancel pending async calls."""
-
-    def __init__(self) -> None:
-        self._cancelled: bool = False
-        self._lock: threading.Lock = threading.Lock()
-        self._callbacks: list[Callable[[], None]] = []
-
-    def cancel(self) -> None:
-        """Cancel pending async calls linked to this cancellation token."""
-        with self._lock:
-            if not self._cancelled:
-                self._cancelled = True
-                for callback in self._callbacks:
-                    callback()
-
-    def is_cancelled(self) -> bool:
-        """Check if the CancellationToken has been used."""
-        with self._lock:
-            return self._cancelled
-
-    def add_callback(self, callback: Callable[[], None]) -> None:
-        """Attach a callback that will be called when cancel is invoked."""
-        with self._lock:
-            if self._cancelled:
-                callback()
-            else:
-                self._callbacks.append(callback)
-
-    def link_future(self, future: Future[Any]) -> Future[Any]:
-        """Link a pending async call to a token to allow its cancellation."""
-        with self._lock:
-            if self._cancelled:
-                future.cancel()
-            else:
-
-                def _cancel() -> None:
-                    future.cancel()
-
-                self._callbacks.append(_cancel)
-        return future
diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py
index 88f7fc7ca0..2cbea3f83f 100644
--- a/python/packages/main/agent_framework/_clients.py
+++ b/python/packages/main/agent_framework/_clients.py
@@ -90,7 +90,7 @@ def _tool_call_non_streaming(
         response: ChatResponse | None = None
         fcc_messages: list[ChatMessage] = []
         for attempt_idx in range(getattr(self, "__maximum_iterations_per_request", 10)):
-            response = await func(self, messages=messages, chat_options=chat_options)
+            response = await func(self, messages=messages, chat_options=chat_options, **kwargs)
             # if there are function calls, we will handle them first
             function_results = {
                 it.call_id for it in response.messages[0].contents if isinstance(it, FunctionResultContent)
diff --git a/python/samples/getting_started/chat_client/chat_response_cancellation.py b/python/samples/getting_started/chat_client/chat_response_cancellation.py
new file mode 100644
index 0000000000..2c1239559e
--- /dev/null
+++ b/python/samples/getting_started/chat_client/chat_response_cancellation.py
@@ -0,0 +1,19 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+
+from agent_framework.openai import OpenAIChatClient
+
+async def main():
+    chat_client = OpenAIChatClient()
+
+    try:
+        task = asyncio.create_task(chat_client.get_response(messages=["Tell me a fantasy story."]))
+        await asyncio.sleep(1)
+        task.cancel()
+        await task
+    except asyncio.CancelledError:
+        print("Request was cancelled")
+
+if __name__ == "__main__":
+    asyncio.run(main())