agno 1.8.1__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

agno/run/{response.py → agent.py}

@@ -7,8 +7,10 @@ from pydantic import BaseModel
 
 from agno.media import AudioArtifact, AudioResponse, ImageArtifact, VideoArtifact
 from agno.models.message import Citations, Message
+from agno.models.metrics import Metrics
 from agno.models.response import ToolExecution
-from agno.run.base import BaseRunResponseEvent, RunResponseExtraData, RunStatus
+from agno.reasoning.step import ReasoningStep
+from agno.run.base import BaseRunOutputEvent, MessageReferences, RunStatus
 from agno.utils.log import logger
 
 
@@ -16,8 +18,8 @@ class RunEvent(str, Enum):
     """Events that can be sent by the run() functions"""
 
     run_started = "RunStarted"
-    run_response_content = "RunResponseContent"
-    run_intermediate_response_content = "RunIntermediateResponseContent"
+    run_content = "RunContent"
+    run_intermediate_content = "RunIntermediateContent"
     run_completed = "RunCompleted"
     run_error = "RunError"
     run_cancelled = "RunCancelled"
@@ -43,14 +45,20 @@ class RunEvent(str, Enum):
 
 
 @dataclass
-class BaseAgentRunResponseEvent(BaseRunResponseEvent):
+class BaseAgentRunEvent(BaseRunOutputEvent):
     created_at: int = field(default_factory=lambda: int(time()))
     event: str = ""
     agent_id: str = ""
     agent_name: str = ""
     run_id: Optional[str] = None
     session_id: Optional[str] = None
-    team_session_id: Optional[str] = None
+
+    # Step context for workflow execution
+    workflow_id: Optional[str] = None
+    workflow_run_id: Optional[str] = None
+    step_id: Optional[str] = None
+    step_name: Optional[str] = None
+    step_index: Optional[int] = None
     tools: Optional[List[ToolExecution]] = None
 
     # For backwards compatibility
@@ -70,7 +78,7 @@ class BaseAgentRunResponseEvent(BaseRunResponseEvent):
 
 
 @dataclass
-class RunResponseStartedEvent(BaseAgentRunResponseEvent):
+class RunStartedEvent(BaseAgentRunEvent):
     """Event sent when the run starts"""
 
     event: str = RunEvent.run_started.value
@@ -79,44 +87,50 @@ class RunResponseStartedEvent(BaseAgentRunResponseEvent):
 
 
 @dataclass
-class RunResponseContentEvent(BaseAgentRunResponseEvent):
-    """Main event for each delta of the RunResponse"""
+class RunContentEvent(BaseAgentRunEvent):
+    """Main event for each delta of the RunOutput"""
 
-    event: str = RunEvent.run_response_content.value
+    event: str = RunEvent.run_content.value
     content: Optional[Any] = None
     content_type: str = "str"
-    thinking: Optional[str] = None
     reasoning_content: Optional[str] = None
     citations: Optional[Citations] = None
     response_audio: Optional[AudioResponse] = None  # Model audio response
     image: Optional[ImageArtifact] = None  # Image attached to the response
-    extra_data: Optional[RunResponseExtraData] = None
+    references: Optional[List[MessageReferences]] = None
+    additional_input: Optional[List[Message]] = None
+    reasoning_steps: Optional[List[ReasoningStep]] = None
+    reasoning_messages: Optional[List[Message]] = None
 
 
 @dataclass
-class IntermediateRunResponseContentEvent(BaseAgentRunResponseEvent):
-    event: str = RunEvent.run_intermediate_response_content.value
+class IntermediateRunContentEvent(BaseAgentRunEvent):
+    event: str = RunEvent.run_intermediate_content.value
     content: Optional[Any] = None
     content_type: str = "str"
 
 
 @dataclass
-class RunResponseCompletedEvent(BaseAgentRunResponseEvent):
+class RunCompletedEvent(BaseAgentRunEvent):
     event: str = RunEvent.run_completed.value
     content: Optional[Any] = None
     content_type: str = "str"
     reasoning_content: Optional[str] = None
-    thinking: Optional[str] = None
     citations: Optional[Citations] = None
     images: Optional[List[ImageArtifact]] = None  # Images attached to the response
     videos: Optional[List[VideoArtifact]] = None  # Videos attached to the response
     audio: Optional[List[AudioArtifact]] = None  # Audio attached to the response
     response_audio: Optional[AudioResponse] = None  # Model audio response
-    extra_data: Optional[RunResponseExtraData] = None
+    references: Optional[List[MessageReferences]] = None
+    additional_input: Optional[List[Message]] = None
+    reasoning_steps: Optional[List[ReasoningStep]] = None
+    reasoning_messages: Optional[List[Message]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    metrics: Optional[Metrics] = None
 
 
 @dataclass
-class RunResponsePausedEvent(BaseAgentRunResponseEvent):
+class RunPausedEvent(BaseAgentRunEvent):
     event: str = RunEvent.run_paused.value
     tools: Optional[List[ToolExecution]] = None
 
@@ -126,18 +140,18 @@ class RunResponsePausedEvent(BaseAgentRunResponseEvent):
 
 
 @dataclass
-class RunResponseContinuedEvent(BaseAgentRunResponseEvent):
+class RunContinuedEvent(BaseAgentRunEvent):
     event: str = RunEvent.run_continued.value
 
 
 @dataclass
-class RunResponseErrorEvent(BaseAgentRunResponseEvent):
+class RunErrorEvent(BaseAgentRunEvent):
     event: str = RunEvent.run_error.value
     content: Optional[str] = None
 
 
 @dataclass
-class RunResponseCancelledEvent(BaseAgentRunResponseEvent):
+class RunCancelledEvent(BaseAgentRunEvent):
     event: str = RunEvent.run_cancelled.value
     reason: Optional[str] = None
 
@@ -147,22 +161,22 @@ class RunResponseCancelledEvent(BaseAgentRunResponseEvent):
 
 
 @dataclass
-class MemoryUpdateStartedEvent(BaseAgentRunResponseEvent):
+class MemoryUpdateStartedEvent(BaseAgentRunEvent):
     event: str = RunEvent.memory_update_started.value
 
 
 @dataclass
-class MemoryUpdateCompletedEvent(BaseAgentRunResponseEvent):
+class MemoryUpdateCompletedEvent(BaseAgentRunEvent):
     event: str = RunEvent.memory_update_completed.value
 
 
 @dataclass
-class ReasoningStartedEvent(BaseAgentRunResponseEvent):
+class ReasoningStartedEvent(BaseAgentRunEvent):
     event: str = RunEvent.reasoning_started.value
 
 
 @dataclass
-class ReasoningStepEvent(BaseAgentRunResponseEvent):
+class ReasoningStepEvent(BaseAgentRunEvent):
     event: str = RunEvent.reasoning_step.value
     content: Optional[Any] = None
     content_type: str = "str"
@@ -170,20 +184,20 @@ class ReasoningStepEvent(BaseAgentRunResponseEvent):
 
 
 @dataclass
-class ReasoningCompletedEvent(BaseAgentRunResponseEvent):
+class ReasoningCompletedEvent(BaseAgentRunEvent):
     event: str = RunEvent.reasoning_completed.value
     content: Optional[Any] = None
     content_type: str = "str"
 
 
 @dataclass
-class ToolCallStartedEvent(BaseAgentRunResponseEvent):
+class ToolCallStartedEvent(BaseAgentRunEvent):
     event: str = RunEvent.tool_call_started.value
     tool: Optional[ToolExecution] = None
 
 
 @dataclass
-class ToolCallCompletedEvent(BaseAgentRunResponseEvent):
+class ToolCallCompletedEvent(BaseAgentRunEvent):
     event: str = RunEvent.tool_call_completed.value
     tool: Optional[ToolExecution] = None
     content: Optional[Any] = None
@@ -193,34 +207,34 @@ class ToolCallCompletedEvent(BaseAgentRunResponseEvent):
 
 
 @dataclass
-class ParserModelResponseStartedEvent(BaseAgentRunResponseEvent):
+class ParserModelResponseStartedEvent(BaseAgentRunEvent):
     event: str = RunEvent.parser_model_response_started.value
 
 
 @dataclass
-class ParserModelResponseCompletedEvent(BaseAgentRunResponseEvent):
+class ParserModelResponseCompletedEvent(BaseAgentRunEvent):
     event: str = RunEvent.parser_model_response_completed.value
 
 
 @dataclass
-class OutputModelResponseStartedEvent(BaseAgentRunResponseEvent):
+class OutputModelResponseStartedEvent(BaseAgentRunEvent):
     event: str = RunEvent.output_model_response_started.value
 
 
 @dataclass
-class OutputModelResponseCompletedEvent(BaseAgentRunResponseEvent):
+class OutputModelResponseCompletedEvent(BaseAgentRunEvent):
     event: str = RunEvent.output_model_response_completed.value
 
 
-RunResponseEvent = Union[
-    RunResponseStartedEvent,
-    RunResponseContentEvent,
-    IntermediateRunResponseContentEvent,
-    RunResponseCompletedEvent,
-    RunResponseErrorEvent,
-    RunResponseCancelledEvent,
-    RunResponsePausedEvent,
-    RunResponseContinuedEvent,
+RunOutputEvent = Union[
+    RunStartedEvent,
+    RunContentEvent,
+    IntermediateRunContentEvent,
+    RunCompletedEvent,
+    RunErrorEvent,
+    RunCancelledEvent,
+    RunPausedEvent,
+    RunContinuedEvent,
     ReasoningStartedEvent,
     ReasoningStepEvent,
     ReasoningCompletedEvent,
@@ -237,14 +251,14 @@ RunResponseEvent = Union[
 
 # Map event string to dataclass
 RUN_EVENT_TYPE_REGISTRY = {
-    RunEvent.run_started.value: RunResponseStartedEvent,
-    RunEvent.run_response_content.value: RunResponseContentEvent,
-    RunEvent.run_intermediate_response_content.value: IntermediateRunResponseContentEvent,
-    RunEvent.run_completed.value: RunResponseCompletedEvent,
-    RunEvent.run_error.value: RunResponseErrorEvent,
-    RunEvent.run_cancelled.value: RunResponseCancelledEvent,
-    RunEvent.run_paused.value: RunResponsePausedEvent,
-    RunEvent.run_continued.value: RunResponseContinuedEvent,
+    RunEvent.run_started.value: RunStartedEvent,
+    RunEvent.run_content.value: RunContentEvent,
+    RunEvent.run_intermediate_content.value: IntermediateRunContentEvent,
+    RunEvent.run_completed.value: RunCompletedEvent,
+    RunEvent.run_error.value: RunErrorEvent,
+    RunEvent.run_cancelled.value: RunCancelledEvent,
+    RunEvent.run_paused.value: RunPausedEvent,
+    RunEvent.run_continued.value: RunContinuedEvent,
     RunEvent.reasoning_started.value: ReasoningStartedEvent,
     RunEvent.reasoning_step.value: ReasoningStepEvent,
     RunEvent.reasoning_completed.value: ReasoningCompletedEvent,
@@ -259,7 +273,7 @@ RUN_EVENT_TYPE_REGISTRY = {
 }
 
 
-def run_response_event_from_dict(data: dict) -> BaseRunResponseEvent:
+def run_output_event_from_dict(data: dict) -> BaseRunOutputEvent:
     event_type = data.get("event", "")
     cls = RUN_EVENT_TYPE_REGISTRY.get(event_type)
     if not cls:
@@ -268,37 +282,53 @@ def run_response_event_from_dict(data: dict) -> BaseRunResponseEvent:
 
 
 @dataclass
-class RunResponse:
+class RunOutput:
     """Response returned by Agent.run() or Workflow.run() functions"""
 
-    content: Optional[Any] = None
-    content_type: str = "str"
-    thinking: Optional[str] = None
-    reasoning_content: Optional[str] = None
-    messages: Optional[List[Message]] = None
-    metrics: Optional[Dict[str, Any]] = None
-    model: Optional[str] = None
-    model_provider: Optional[str] = None
     run_id: Optional[str] = None
     agent_id: Optional[str] = None
     agent_name: Optional[str] = None
     session_id: Optional[str] = None
-    team_session_id: Optional[str] = None
+    parent_run_id: Optional[str] = None
     workflow_id: Optional[str] = None
+    user_id: Optional[str] = None
+
+    content: Optional[Any] = None
+    content_type: str = "str"
+
+    reasoning_content: Optional[str] = None
+    reasoning_steps: Optional[List[ReasoningStep]] = None
+    reasoning_messages: Optional[List[Message]] = None
+
+    model: Optional[str] = None
+    model_provider: Optional[str] = None
+    messages: Optional[List[Message]] = None
+    metrics: Optional[Metrics] = None
+    additional_input: Optional[List[Message]] = None
+
     tools: Optional[List[ToolExecution]] = None
-    formatted_tool_calls: Optional[List[str]] = None
+
     images: Optional[List[ImageArtifact]] = None  # Images attached to the response
     videos: Optional[List[VideoArtifact]] = None  # Videos attached to the response
     audio: Optional[List[AudioArtifact]] = None  # Audio attached to the response
     response_audio: Optional[AudioResponse] = None  # Model audio response
+
     citations: Optional[Citations] = None
-    extra_data: Optional[RunResponseExtraData] = None
+    references: Optional[List[MessageReferences]] = None
+
+    metadata: Optional[Dict[str, Any]] = None
+
     created_at: int = field(default_factory=lambda: int(time()))
 
-    events: Optional[List[RunResponseEvent]] = None
+    events: Optional[List[RunOutputEvent]] = None
 
     status: RunStatus = RunStatus.running
 
+    # === FOREIGN KEY RELATIONSHIPS ===
+    # These fields establish relationships to parent workflow/step structures
+    # and should be treated as foreign keys for data integrity
+    workflow_step_id: Optional[str] = None  # FK: Points to StepOutput.step_id
+
     @property
     def is_paused(self):
         return self.status == RunStatus.paused
@@ -328,16 +358,23 @@ class RunResponse:
             not in [
                 "messages",
                 "tools",
-                "extra_data",
+                "metadata",
                 "images",
                 "videos",
                 "audio",
                 "response_audio",
                 "citations",
                 "events",
+                "additional_input",
+                "reasoning_steps",
+                "reasoning_messages",
+                "references",
             ]
         }
 
+        if self.metrics is not None:
+            _dict["metrics"] = self.metrics.to_dict() if isinstance(self.metrics, Metrics) else self.metrics
+
         if self.events is not None:
             _dict["events"] = [e.to_dict() for e in self.events]
 
@@ -347,10 +384,20 @@ class RunResponse:
         if self.messages is not None:
             _dict["messages"] = [m.to_dict() for m in self.messages]
 
-        if self.extra_data is not None:
-            _dict["extra_data"] = (
-                self.extra_data.to_dict() if isinstance(self.extra_data, RunResponseExtraData) else self.extra_data
-            )
+        if self.metadata is not None:
+            _dict["metadata"] = self.metadata
+
+        if self.additional_input is not None:
+            _dict["additional_input"] = [m.to_dict() for m in self.additional_input]
+
+        if self.reasoning_messages is not None:
+            _dict["reasoning_messages"] = [m.to_dict() for m in self.reasoning_messages]
+
+        if self.reasoning_steps is not None:
+            _dict["reasoning_steps"] = [rs.model_dump() for rs in self.reasoning_steps]
+
+        if self.references is not None:
+            _dict["references"] = [r.model_dump() for r in self.references]
 
         if self.images is not None:
             _dict["images"] = []
@@ -413,9 +460,12 @@ class RunResponse:
         return json.dumps(_dict, indent=2)
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "RunResponse":
+    def from_dict(cls, data: Dict[str, Any]) -> "RunOutput":
+        if "run" in data:
+            data = data.pop("run")
+
         events = data.pop("events", None)
-        events = [run_response_event_from_dict(event) for event in events] if events else None
+        events = [run_output_event_from_dict(event) for event in events] if events else None
 
         messages = data.pop("messages", None)
         messages = [Message.model_validate(message) for message in messages] if messages else None
@@ -438,12 +488,30 @@ class RunResponse:
         response_audio = data.pop("response_audio", None)
         response_audio = AudioResponse.model_validate(response_audio) if response_audio else None
 
-        # To make it backwards compatible
-        if "event" in data:
-            data.pop("event")
+        metrics = data.pop("metrics", None)
+        if metrics:
+            metrics = Metrics(**metrics)
+
+        additional_input = data.pop("additional_input", None)
+
+        if additional_input is not None:
+            additional_input = [Message.model_validate(message) for message in additional_input]
+
+        reasoning_steps = data.pop("reasoning_steps", None)
+        if reasoning_steps is not None:
+            reasoning_steps = [ReasoningStep.model_validate(step) for step in reasoning_steps]
+
+        reasoning_messages = data.pop("reasoning_messages", None)
+        if reasoning_messages is not None:
+            reasoning_messages = [Message.model_validate(message) for message in reasoning_messages]
+
+        references = data.pop("references", None)
+        if references is not None:
+            references = [MessageReferences.model_validate(reference) for reference in references]
 
         return cls(
             messages=messages,
+            metrics=metrics,
             citations=citations,
             tools=tools,
             images=images,
@@ -451,6 +519,10 @@ class RunResponse:
             videos=videos,
             response_audio=response_audio,
             events=events,
+            additional_input=additional_input,
+            reasoning_steps=reasoning_steps,
+            reasoning_messages=reasoning_messages,
+            references=references,
             **data,
         )
 

agno/run/base.py

@@ -1,18 +1,19 @@
 from dataclasses import asdict, dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict
 
 from pydantic import BaseModel
 
 from agno.media import AudioArtifact, AudioResponse, ImageArtifact, VideoArtifact
 from agno.models.message import Citations, Message, MessageReferences
+from agno.models.metrics import Metrics
 from agno.models.response import ToolExecution
 from agno.reasoning.step import ReasoningStep
 from agno.utils.log import log_error
 
 
 @dataclass
-class BaseRunResponseEvent:
+class BaseRunOutputEvent:
     def to_dict(self) -> Dict[str, Any]:
         _dict = {
             k: v
@@ -22,7 +23,7 @@ class BaseRunResponseEvent:
             not in [
                 "tools",
                 "tool",
-                "extra_data",
+                "metadata",
                 "image",
                 "images",
                 "videos",
@@ -30,13 +31,28 @@ class BaseRunResponseEvent:
                 "response_audio",
                 "citations",
                 "member_responses",
+                "reasoning_messages",
+                "reasoning_steps",
+                "references",
+                "additional_input",
+                "metrics",
             ]
         }
 
-        if hasattr(self, "extra_data") and self.extra_data is not None:
-            _dict["extra_data"] = (
-                self.extra_data.to_dict() if isinstance(self.extra_data, RunResponseExtraData) else self.extra_data
-            )
+        if hasattr(self, "metadata") and self.metadata is not None:
+            _dict["metadata"] = self.metadata
+
+        if hasattr(self, "additional_input") and self.additional_input is not None:
+            _dict["additional_input"] = [m.to_dict() for m in self.additional_input]
+
+        if hasattr(self, "reasoning_messages") and self.reasoning_messages is not None:
+            _dict["reasoning_messages"] = [m.to_dict() for m in self.reasoning_messages]
+
+        if hasattr(self, "reasoning_steps") and self.reasoning_steps is not None:
+            _dict["reasoning_steps"] = [rs.model_dump() for rs in self.reasoning_steps]
+
+        if hasattr(self, "references") and self.references is not None:
+            _dict["references"] = [r.model_dump() for r in self.references]
 
         if hasattr(self, "member_responses") and self.member_responses:
             _dict["member_responses"] = [response.to_dict() for response in self.member_responses]
@@ -100,6 +116,9 @@ class BaseRunResponseEvent:
             else:
                 _dict["tool"] = self.tool
 
+        if hasattr(self, "metrics") and self.metrics is not None:
+            _dict["metrics"] = self.metrics.to_dict()
+
         return _dict
 
     def to_json(self) -> str:
@@ -139,13 +158,25 @@ class BaseRunResponseEvent:
         if response_audio:
             data["response_audio"] = AudioResponse.model_validate(response_audio)
 
-        extra_data = data.pop("extra_data", None)
-        if extra_data:
-            data["extra_data"] = RunResponseExtraData.from_dict(extra_data)
+        additional_input = data.pop("additional_input", None)
+        if additional_input is not None:
+            data["additional_input"] = [Message.model_validate(message) for message in additional_input]
 
-        # To make it backwards compatible
-        if "event" in data:
-            data.pop("event")
+        reasoning_steps = data.pop("reasoning_steps", None)
+        if reasoning_steps is not None:
+            data["reasoning_steps"] = [ReasoningStep.model_validate(step) for step in reasoning_steps]
+
+        reasoning_messages = data.pop("reasoning_messages", None)
+        if reasoning_messages is not None:
+            data["reasoning_messages"] = [Message.model_validate(message) for message in reasoning_messages]
+
+        references = data.pop("references", None)
+        if references is not None:
+            data["references"] = [MessageReferences.model_validate(reference) for reference in references]
+
+        metrics = data.pop("metrics", None)
+        if metrics:
+            data["metrics"] = Metrics(**metrics)
 
         return cls(**data)
 
@@ -158,51 +189,6 @@ class BaseRunResponseEvent:
         return False
 
 
-@dataclass
-class RunResponseExtraData:
-    references: Optional[List[MessageReferences]] = None
-    add_messages: Optional[List[Message]] = None
-    reasoning_steps: Optional[List[ReasoningStep]] = None
-    reasoning_messages: Optional[List[Message]] = None
-
-    def to_dict(self) -> Dict[str, Any]:
-        _dict = {}
-        if self.add_messages is not None:
-            _dict["add_messages"] = [m.to_dict() for m in self.add_messages]
-        if self.reasoning_messages is not None:
-            _dict["reasoning_messages"] = [m.to_dict() for m in self.reasoning_messages]
-        if self.reasoning_steps is not None:
-            _dict["reasoning_steps"] = [rs.model_dump() for rs in self.reasoning_steps]
-        if self.references is not None:
-            _dict["references"] = [r.model_dump() for r in self.references]
-        return _dict
-
-    @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "RunResponseExtraData":
-        add_messages = data.pop("add_messages", None)
-        if add_messages is not None:
-            add_messages = [Message.model_validate(message) for message in add_messages]
-
-        reasoning_steps = data.pop("reasoning_steps", None)
-        if reasoning_steps is not None:
-            reasoning_steps = [ReasoningStep.model_validate(step) for step in reasoning_steps]
-
-        reasoning_messages = data.pop("reasoning_messages", None)
-        if reasoning_messages is not None:
-            reasoning_messages = [Message.model_validate(message) for message in reasoning_messages]
-
-        references = data.pop("references", None)
-        if references is not None:
-            references = [MessageReferences.model_validate(reference) for reference in references]
-
-        return cls(
-            add_messages=add_messages,
-            reasoning_steps=reasoning_steps,
-            reasoning_messages=reasoning_messages,
-            references=references,
-        )
-
-
 class RunStatus(str, Enum):
     """State of the main run response"""
 

agno/run/cancel.py

@@ -0,0 +1,83 @@
+"""Run cancellation management."""
+
+import threading
+from typing import Dict
+
+from agno.exceptions import RunCancelledException
+from agno.utils.log import logger
+
+
+class RunCancellationManager:
+    """Manages cancellation state for agent runs."""
+
+    def __init__(self):
+        self._cancelled_runs: Dict[str, bool] = {}
+        self._lock = threading.Lock()
+
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        with self._lock:
+            self._cancelled_runs[run_id] = False
+            logger.debug(f"Registered run {run_id} for cancellation tracking")
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                self._cancelled_runs[run_id] = True
+                logger.info(f"Run {run_id} marked for cancellation")
+                return True
+            else:
+                logger.warning(f"Attempted to cancel unknown run {run_id}")
+                return False
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        with self._lock:
+            return self._cancelled_runs.get(run_id, False)
+
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                del self._cancelled_runs[run_id]
+                logger.debug(f"Cleaned up cancellation tracking for run {run_id}")
+
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        if self.is_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status."""
+        with self._lock:
+            return self._cancelled_runs.copy()
+
+
+# Global cancellation manager instance
+_cancellation_manager = RunCancellationManager()
+
+
+def register_run(run_id: str) -> None:
+    """Register a new run for cancellation tracking."""
+    _cancellation_manager.register_run(run_id)
+
+
+def cancel_run(run_id: str) -> bool:
+    """Cancel a run."""
+    return _cancellation_manager.cancel_run(run_id)
+
+
+def cleanup_run(run_id: str) -> None:
+    """Clean up cancellation tracking for a completed run."""
+    _cancellation_manager.cleanup_run(run_id)
+
+
+def raise_if_cancelled(run_id: str) -> None:
+    """Check if a run should be cancelled and raise exception if so."""
+    _cancellation_manager.raise_if_cancelled(run_id)