@@ -118,6 +118,8 @@ def __init__(
118118 self .connectors_api = connectors_api
119119 self ._background_queue : asyncio .Queue = asyncio .Queue (maxsize = BACKGROUND_QUEUE_MAX_SIZE )
120120 self ._background_worker_tasks : set [asyncio .Task ] = set ()
121+ self ._background_response_tasks : dict [str , asyncio .Task ] = {}
122+ self ._background_response_tasks_lock = asyncio .Lock ()
121123
122124 async def initialize (self ) -> None :
123125 """No-op: background workers are started lazily on first use.
@@ -137,21 +139,51 @@ async def _ensure_workers_started(self) -> None:
137139
138140 async def shutdown (self ) -> None :
139141 """Stop background worker pool."""
142+ # Cancel all in-progress response tasks
143+ async with self ._background_response_tasks_lock :
144+ for task in self ._background_response_tasks .values ():
145+ task .cancel ()
146+ response_task_list = list (self ._background_response_tasks .values ())
147+
148+ # Cancel worker tasks
140149 for task in self ._background_worker_tasks :
141150 task .cancel ()
142- await asyncio .gather (* self ._background_worker_tasks , return_exceptions = True )
151+
152+ # Wait for all tasks to complete
153+ all_tasks = list (self ._background_worker_tasks ) + response_task_list
154+ await asyncio .gather (* all_tasks , return_exceptions = True )
143155
144156 async def _background_worker (self ) -> None :
145157 """Worker coroutine that pulls items from the queue and processes them."""
146158 while True :
147159 kwargs = await self ._background_queue .get ()
148- try :
149- await asyncio .wait_for (
160+ response_id = kwargs ["response_id" ]
161+
162+ # Create a task for this specific response so we can cancel it
163+ processing_task = asyncio .create_task (
164+ asyncio .wait_for (
150165 self ._run_background_response_loop (** kwargs ),
151166 timeout = BACKGROUND_RESPONSE_TIMEOUT_SECONDS ,
152167 )
168+ )
169+
170+ # Track the task
171+ async with self ._background_response_tasks_lock :
172+ self ._background_response_tasks [response_id ] = processing_task
173+
174+ try :
175+ await processing_task
176+ except asyncio .CancelledError :
177+ # Response was cancelled via cancel_openai_response
178+ logger .info (f"Background response { response_id } )
179+ try :
180+ existing = await self .responses_store .get_response_object (response_id )
181+ if existing .status != "cancelled" :
182+ existing .status = "cancelled"
183+ await self .responses_store .update_response_object (existing )
184+ except Exception :
185+ logger .exception (f"Failed to update response { response_id } )
153186 except TimeoutError :
154- response_id = kwargs ["response_id" ]
155187 logger .exception (
156188 f"Background response { response_id } { BACKGROUND_RESPONSE_TIMEOUT_SECONDS }
157189 )
@@ -169,7 +201,6 @@ async def _background_worker(self) -> None:
169201 "Client polling this response will not see the failure."
170202 )
171203 except Exception as e :
172- response_id = kwargs ["response_id" ]
173204 logger .exception (f"Error processing background response { response_id } )
174205 try :
175206 existing = await self .responses_store .get_response_object (response_id )
@@ -185,6 +216,9 @@ async def _background_worker(self) -> None:
185216 "Client polling this response will not see the failure."
186217 )
187218 finally :
219+ # Remove from tracking
220+ async with self ._background_response_tasks_lock :
221+ self ._background_response_tasks .pop (response_id , None )
188222 self ._background_queue .task_done ()
189223
190224 async def _prepend_previous_response (
@@ -1095,6 +1129,61 @@ async def _create_streaming_response(
10951129 async def delete_openai_response (self , response_id : str ) -> OpenAIDeleteResponseObject :
10961130 return await self .responses_store .delete_response_object (response_id )
10971131
1132+ async def cancel_openai_response (
1133+ self ,
1134+ response_id : str ,
1135+ ) -> OpenAIResponseObject :
1136+ """Cancel a response that is queued or in progress.
1137+
1138+ Args:
1139+ response_id: The ID of the response to cancel
1140+
1141+ Returns:
1142+ The updated response object with status "cancelled"
1143+
1144+ Raises:
1145+ ResponseNotFoundError: If the response doesn't exist (automatically from store)
1146+ ConflictError: If the response is already in a terminal state
1147+ """
1148+ from llama_stack_api import ConflictError
1149+
1150+ # Get current response state
1151+ response = await self .responses_store .get_response_object (response_id )
1152+
1153+ # Preserve the background field
1154+ was_background = response .background
1155+
1156+ # If already cancelled, return current state (idempotent)
1157+ if response .status == "cancelled" :
1158+ cancelled_response = response .to_response_object ()
1159+ # Ensure background field is preserved
1160+ if was_background :
1161+ cancelled_response .background = was_background
1162+ return cancelled_response
1163+
1164+ # Cannot cancel responses in terminal states
1165+ if response .status in ["completed" , "failed" , "incomplete" ]:
1166+ raise ConflictError (f"Cannot cancel response '{ response_id } { response .status } )
1167+
1168+ # Update status to cancelled in database
1169+ response .status = "cancelled"
1170+ await self .responses_store .update_response_object (response )
1171+
1172+ # If the response is currently being processed, cancel the task
1173+ async with self ._background_response_tasks_lock :
1174+ task = self ._background_response_tasks .get (response_id )
1175+ if task :
1176+ task .cancel ()
1177+ # Note: task removal handled in worker's finally block
1178+
1179+ # Return updated response
1180+ response_with_input = await self .responses_store .get_response_object (response_id )
1181+ cancelled_response = response_with_input .to_response_object ()
1182+ # Ensure background field is preserved
1183+ if was_background :
1184+ cancelled_response .background = was_background
1185+ return cancelled_response
1186+
10981187 async def _sync_response_to_conversation (
10991188 self , conversation_id : str , input : str | list [OpenAIResponseInput ] | None , output_items : list [ConversationItem ]
11001189 ) -> None :
0 commit comments