feat(run): add cooperative cancellation, inject, and status event handling

- Added cooperative cancel support to RunResultStreaming - Introduced flag for backward compatibility - Differentiated vs run statuses - Updated internal run loop to handle injected items mid-run - Added initial tests (test_cancel_streamed_run.py) for cancellation behavior

Author: vrtnis

src/agents/_run_impl.py

@@ -3,6 +3,7 @@
 import asyncio
 import dataclasses
 import inspect
+import contextlib
 from collections.abc import Awaitable
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, cast
@@ -225,6 +226,25 @@ def get_model_tracing_impl(
  else:
  return ModelTracing.ENABLED_WITHOUT_DATA
 
+# --- NEW: helpers for cancellable tool execution ---
+
+async def _await_cancellable(awaitable):
+ """Await an awaitable in its own task so CancelledError interrupts promptly."""
+ task = asyncio.create_task(awaitable)
+ try:
+ return await task
+ except asyncio.CancelledError:
+ # propagate so run.py can handle terminal cancel
+ raise
+
+def _maybe_call_cancel_hook(tool_obj) -> None:
+ """Best-effort: call a cancel/terminate hook on the tool if present."""
+ for name in ("cancel", "terminate", "stop"):
+ cb = getattr(tool_obj, name, None)
+ if callable(cb):
+ with contextlib.suppress(Exception):
+ cb()
+ break
 
 class RunImpl:
  @classmethod
@@ -556,24 +576,26 @@ async def run_single_tool(
  if config.trace_include_sensitive_data:
  span_fn.span_data.input = tool_call.arguments
  try:
- _, _, result = await asyncio.gather(
- hooks.on_tool_start(tool_context, agent, func_tool),
- (
- agent.hooks.on_tool_start(tool_context, agent, func_tool)
- if agent.hooks
- else _coro.noop_coroutine()
- ),
- func_tool.on_invoke_tool(tool_context, tool_call.arguments),
- )
+ # run start hooks first (don’t tie them to the cancellable task)
+ await asyncio.gather(
+ hooks.on_tool_start(tool_context, agent, func_tool),
+ (agent.hooks.on_tool_start(tool_context, agent, func_tool) if agent.hooks else _coro.noop_coroutine()),
+ )
+ 
+ try:
+ result = await _await_cancellable(
+ func_tool.on_invoke_tool(tool_context, tool_call.arguments)
+ )
+ except asyncio.CancelledError:
+ _maybe_call_cancel_hook(func_tool)
+ raise
+ 
+ await asyncio.gather(
+ hooks.on_tool_end(tool_context, agent, func_tool, result),
+ (agent.hooks.on_tool_end(tool_context, agent, func_tool, result) if agent.hooks else _coro.noop_coroutine()),
+ )
+ 
 
- await asyncio.gather(
- hooks.on_tool_end(tool_context, agent, func_tool, result),
- (
- agent.hooks.on_tool_end(tool_context, agent, func_tool, result)
- if agent.hooks
- else _coro.noop_coroutine()
- ),
- )
  except Exception as e:
  _error_tracing.attach_error_to_current_span(
  SpanError(
@@ -643,44 +665,45 @@ async def execute_computer_actions(
  context_wrapper: RunContextWrapper[TContext],
  config: RunConfig,
  ) -> list[RunItem]:
- results: list[RunItem] = []
- # Need to run these serially, because each action can affect the computer state
- for action in actions:
- acknowledged: list[ComputerCallOutputAcknowledgedSafetyCheck] | None = None
- if action.tool_call.pending_safety_checks and action.computer_tool.on_safety_check:
- acknowledged = []
- for check in action.tool_call.pending_safety_checks:
- data = ComputerToolSafetyCheckData(
- ctx_wrapper=context_wrapper,
- agent=agent,
- tool_call=action.tool_call,
- safety_check=check,
- )
- maybe = action.computer_tool.on_safety_check(data)
- ack = await maybe if inspect.isawaitable(maybe) else maybe
- if ack:
- acknowledged.append(
- ComputerCallOutputAcknowledgedSafetyCheck(
- id=check.id,
- code=check.code,
- message=check.message,
- )
- )
- else:
- raise UserError("Computer tool safety check was not acknowledged")
-
- results.append(
- await ComputerAction.execute(
- agent=agent,
- action=action,
- hooks=hooks,
- context_wrapper=context_wrapper,
- config=config,
- acknowledged_safety_checks=acknowledged,
- )
- )
-
- return results
+ results: list[RunItem] = []
+ for action in actions:
+ acknowledged: list[ComputerCallOutputAcknowledgedSafetyCheck] | None = None
+ if action.tool_call.pending_safety_checks and action.computer_tool.on_safety_check:
+ acknowledged = []
+ for check in action.tool_call.pending_safety_checks:
+ data = ComputerToolSafetyCheckData(
+ ctx_wrapper=context_wrapper,
+ agent=agent,
+ tool_call=action.tool_call,
+ safety_check=check,
+ )
+ maybe = action.computer_tool.on_safety_check(data)
+ ack = await maybe if inspect.isawaitable(maybe) else maybe
+ if ack:
+ acknowledged.append(ComputerCallOutputAcknowledgedSafetyCheck(
+ id=check.id, code=check.code, message=check.message
+ ))
+ else:
+ raise UserError("Computer tool safety check was not acknowledged")
+ 
+ try:
+ item = await _await_cancellable(
+ ComputerAction.execute(
+ agent=agent,
+ action=action,
+ hooks=hooks,
+ context_wrapper=context_wrapper,
+ config=config,
+ acknowledged_safety_checks=acknowledged,
+ )
+ )
+ except asyncio.CancelledError:
+ _maybe_call_cancel_hook(action.computer_tool)
+ raise
+ 
+ results.append(item)
+ 
+ return results
 
  @classmethod
  async def execute_handoffs(
@@ -1052,16 +1075,23 @@ async def execute(
  else cls._get_screenshot_sync(action.computer_tool.computer, action.tool_call)
  )
 
- _, _, output = await asyncio.gather(
+ # start hooks first
+ await asyncio.gather(
  hooks.on_tool_start(context_wrapper, agent, action.computer_tool),
  (
  agent.hooks.on_tool_start(context_wrapper, agent, action.computer_tool)
  if agent.hooks
  else _coro.noop_coroutine()
  ),
- output_func,
  )
-
+ # run the action (screenshot/etc) in a cancellable task
+ try:
+ output = await _await_cancellable(output_func)
+ except asyncio.CancelledError:
+ _maybe_call_cancel_hook(action.computer_tool)
+ raise
+ 
+ # end hooks
  await asyncio.gather(
  hooks.on_tool_end(context_wrapper, agent, action.computer_tool, output),
  (
@@ -1169,10 +1199,20 @@ async def execute(
  data=call.tool_call,
  )
  output = call.local_shell_tool.executor(request)
- if inspect.isawaitable(output):
- result = await output
- else:
- result = output
+ try:
+ if inspect.isawaitable(output):
+ result = await _await_cancellable(output)
+ else:
+ # If executor returns a sync result, just use it (can’t cancel mid-call)
+ result = output
+ except asyncio.CancelledError:
+ # Best-effort: if the executor or tool exposes a cancel/terminate / kill, call it
+ _maybe_call_cancel_hook(call.local_shell_tool)
+ # If your executor returns a proc handle (common pattern), adddress it here if needed:
+ # with contextlib.suppress(Exception):
+ # proc.terminate(); await asyncio.wait_for(proc.wait(), 1.0)
+ # proc.kill()
+ raise
 
  await asyncio.gather(
  hooks.on_tool_end(context_wrapper, agent, call.local_shell_tool, result),
@@ -1185,7 +1225,7 @@ async def execute(
 
  return ToolCallOutputItem(
  agent=agent,
- output=output,
+ output=result,
  raw_item={
  "type": "local_shell_call_output",
  "id": call.tool_call.call_id,

src/agents/models/openai_responses.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import asyncio
 from collections.abc import AsyncIterator
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Literal, cast, overload
@@ -171,16 +172,31 @@ async def stream_response(
  )
 
  final_response: Response | None = None
-
- async for chunk in stream:
- if isinstance(chunk, ResponseCompletedEvent):
- final_response = chunk.response
- yield chunk
-
+ 
+ try:
+ async for chunk in stream: # type: ignore[arg-type] # ensure type checkers relax here
+ if isinstance(chunk, ResponseCompletedEvent):
+ final_response = chunk.response
+ yield chunk
+ except asyncio.CancelledError:
+ # Cooperative cancel: ensure the HTTP stream is closed, then propagate
+ try:
+ await stream.aclose()
+ except Exception:
+ pass
+ raise
+ finally:
+ # Always close the stream if the async iterator exits (normal or error)
+ try:
+ await stream.aclose()
+ except Exception:
+ pass
+ 
  if final_response and tracing.include_data():
  span_response.span_data.response = final_response
  span_response.span_data.input = input
-
+ 
+ 
  except Exception as e:
  span_response.set_error(
  SpanError(

src/agents/result.py

@@ -2,6 +2,7 @@
 
 import abc
 import asyncio
+import contextlib
 from collections.abc import AsyncIterator
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, cast
@@ -143,6 +144,10 @@ class RunResultStreaming(RunResultBase):
  is_complete: bool = False
  """Whether the agent has finished running."""
 
+ _emit_status_events: bool = False
+ """Whether to emit RunUpdatedStreamEvent status updates (default False for backward compatibility)."""
+
+
  # Queues that the background run_loop writes to
  _event_queue: asyncio.Queue[StreamEvent | QueueCompleteSentinel] = field(
  default_factory=asyncio.Queue, repr=False
@@ -164,17 +169,53 @@ def last_agent(self) -> Agent[Any]:
  """
  return self.current_agent
 
- def cancel(self) -> None:
- """Cancels the streaming run, stopping all background tasks and marking the run as
- complete."""
- self._cleanup_tasks() # Cancel all running tasks
- self.is_complete = True # Mark the run as complete to stop event streaming
-
- # Optionally, clear the event queue to prevent processing stale events
- while not self._event_queue.empty():
- self._event_queue.get_nowait()
- while not self._input_guardrail_queue.empty():
- self._input_guardrail_queue.get_nowait()
+ 
+ def cancel(self, reason: str | None = None) -> None:
+ # 1) Signal cooperative cancel to the runner
+ active = getattr(self, "_active_run", None)
+ if active:
+ with contextlib.suppress(Exception):
+ active.cancel(reason)
+ 
+ # 2) Wake any stream_events() consumer immediately
+ with contextlib.suppress(Exception):
+ self._event_queue.put_nowait(QueueCompleteSentinel())
+ 
+ # 3) Do NOT cancel the background task; let the loop unwind cooperatively
+ # task = getattr(self, "_run_impl_task", None)
+ # if task and not task.done():
+ # with contextlib.suppress(Exception):
+ # task.cancel()
+ 
+ # 4) Mark complete; flushing only when status events are disabled
+ self.is_complete = True
+ if not getattr(self, "_emit_status_events", False):
+ with contextlib.suppress(Exception):
+ while not self._event_queue.empty():
+ self._event_queue.get_nowait()
+ self._event_queue.task_done()
+ with contextlib.suppress(Exception):
+ while not self._input_guardrail_queue.empty():
+ self._input_guardrail_queue.get_nowait()
+ self._input_guardrail_queue.task_done()
+
+
+ def inject(self, items: list[TResponseInputItem]) -> None:
+ """
+ Inject new input items mid-run. They will be consumed at the start of the next step.
+ """
+ active = getattr(self, "_active_run", None)
+ if active is not None:
+ try:
+ active.inject(items)
+ except Exception:
+ pass
+ 
+ @property
+ def active_run(self):
+ """Access the underlying ActiveRun handle (may be None early in startup)."""
+ return getattr(self, "_active_run", None)
+ 
 
  async def stream_events(self) -> AsyncIterator[StreamEvent]:
  """Stream deltas for new items as they are generated. We're using the types from the